MCPcopy Create free account
hub / github.com/Luce-Org/lucebox-hub / build_draft_graph_internal

Function build_draft_graph_internal

server/src/common/dflash_draft_graph.cpp:31–136  ·  view source on GitHub ↗

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. ctx_alloc: allocation/topology size of the ctx dimension (>= ctx_len). When ctx_alloc > 0 and differs from the legacy behavior, a full-layer pad mask input is created so the graph topology stays stable while ctx_len grows (CUDA-graph replay for the draft forwa

Source from the content-addressed store, hash-verified

29// mask input is created so the graph topology stays stable while ctx_len
30// grows (CUDA-graph replay for the draft forward).
31static bool build_draft_graph_internal(
32 StepGraph & sg,
33 const DraftWeights & dw,
34 ggml_tensor * lm_head,
35 int ctx_len,
36 const DraftFeatureMirror * mirror,
37 int mirror_slot0,
38 bool mirror_view,
39 bool pad_masked = false) {
40
41 const size_t arena_sz = 32u * 1024 * 1024;
42 if (sg.meta_arena.size() < arena_sz) sg.meta_arena.resize(arena_sz);
43 ggml_init_params ip{};
44 ip.mem_size = sg.meta_arena.size();
45 ip.mem_buffer = sg.meta_arena.data();
46 ip.no_alloc = true;
47 sg.ctx = ggml_init(ip);
48 if (!sg.ctx) return false;
49
50 const int hidden = dw.n_embd;
51 const int q_len = dw.block_size;
52 const int fc_in = dw.n_target_layers * hidden;
53
54 sg.inp_embed = ggml_new_tensor_3d(sg.ctx, GGML_TYPE_F32, hidden, q_len, 1);
55 ggml_set_name(sg.inp_embed, "inp_embed");
56 ggml_set_input(sg.inp_embed);
57
58 if (mirror_view) {
59 const size_t stride = mirror->target_feat->nb[1];
60 sg.target_hidden_cat = ggml_view_3d(
61 sg.ctx,
62 mirror->target_feat,
63 fc_in, ctx_len, 1,
64 stride,
65 stride * (size_t)ctx_len,
66 (size_t)mirror_slot0 * stride);
67 } else {
68 sg.target_hidden_cat = ggml_new_tensor_3d(sg.ctx, GGML_TYPE_F32, fc_in, ctx_len, 1);
69 ggml_set_input(sg.target_hidden_cat);
70 }
71 ggml_set_name(sg.target_hidden_cat, "target_hidden_cat");
72
73 sg.positions = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, q_len);
74 ggml_set_name(sg.positions, "positions_q");
75 ggml_set_input(sg.positions);
76
77 sg.positions_k = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, ctx_len + q_len);
78 ggml_set_name(sg.positions_k, "positions_k");
79 ggml_set_input(sg.positions_k);
80
81 // Causal mask for SWA layers (if any).
82 // Shape: [kv_pad, q_len] F16 (directly, no cast needed — matches attn_masks.h pattern).
83 sg.attn_mask = nullptr;
84 const bool has_swa = draft_has_swa_layers(dw);
85 if (has_swa) {
86 // SWA layers' effective KV length (windowed or full ctx)
87 const bool swa_active = dw.swa_window > 0 && ctx_len > dw.swa_window;
88 const int eff_ctx = swa_active ? dw.swa_window : ctx_len;

Callers 1

build_draft_stepFunction · 0.85

Calls 5

draft_has_swa_layersFunction · 0.85
mask_align_upFunction · 0.85
build_draft_graphFunction · 0.85
sizeMethod · 0.45
dataMethod · 0.45

Tested by

no test coverage detected