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

Function build_draft_step

server/src/common/dflash_draft_graph.cpp:138–279  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

136}
137
138bool build_draft_step(
139 StepGraph & sg,
140 const DraftWeights & dw,
141 ggml_tensor * lm_head,
142 ggml_backend_t backend,
143 int ctx_len,
144 const DraftFeatureMirror * mirror,
145 int committed,
146 int /*ctx_len_max*/,
147 bool pad_ctx) {
148 // [TAG_FUSED_LOOP] persistent-graph fast path: when the previous build
149 // used the padded COPY mode (no mirror view baked into the topology) and
150 // ctx_len still lands in the same 64-aligned bucket, every tensor address
151 // and the graph topology are identical — skip the rebuild and only
152 // refresh the two ctx_len-dependent masks. Feature contents, positions
153 // and noise embeds are re-uploaded by the caller every step regardless,
154 // and the pad feature rows beyond ctx_len were zeroed at bucket build.
155 if (sg.gf && sg.ctx_alloc > 0 && !sg.built_view && !mirror && pad_ctx &&
156 sg.ctx_alloc == ((ctx_len + 63) & ~63)) {
157 const int q_len = dw.block_size;
158 const int ctx_alloc = sg.ctx_alloc;
159 static constexpr uint16_t ZERO = 0x0000;
160 static constexpr uint16_t NEG_INF = 0xFC00;
161 if (sg.pad_mask_full) {
162 const int kv_pad = mask_align_up(ctx_alloc + q_len, MASK_KV_PAD);
163 std::vector<uint16_t> mask_data((size_t)kv_pad * q_len, NEG_INF);
164 for (int q = 0; q < q_len; q++) {
165 for (int k = 0; k < ctx_len; k++)
166 mask_data[(size_t)q * kv_pad + k] = ZERO;
167 for (int j = 0; j < q_len; j++)
168 mask_data[(size_t)q * kv_pad + (ctx_alloc + j)] = ZERO;
169 }
170 ggml_backend_tensor_set(sg.pad_mask_full, mask_data.data(), 0,
171 sizeof(uint16_t) * mask_data.size());
172 }
173 if (sg.attn_mask) {
174 const int kv_pad = mask_align_up(ctx_alloc + q_len, MASK_KV_PAD);
175 std::vector<uint16_t> mask_data((size_t)kv_pad * q_len, NEG_INF);
176 for (int q = 0; q < q_len; q++) {
177 for (int k = 0; k < ctx_len; k++)
178 mask_data[(size_t)q * kv_pad + k] = ZERO;
179 for (int j = 0; j <= q; j++)
180 mask_data[(size_t)q * kv_pad + (ctx_alloc + j)] = ZERO;
181 }
182 ggml_backend_tensor_set(sg.attn_mask, mask_data.data(), 0,
183 sizeof(uint16_t) * mask_data.size());
184 }
185 return true;
186 }
187
188 step_graph_free(sg);
189
190 if (!sg.alloc) {
191 sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
192 }
193
194 // Padded-ctx mode: build the graph at the 64-aligned ctx size and mask the
195 // pad keys, so the topology (and gallocr layout) stays IDENTICAL across

Callers 8

mainFunction · 0.85
do_hybrid_spec_decodeMethod · 0.85
run_dflash_spec_decodeFunction · 0.85
do_spec_decodeMethod · 0.85
do_spec_decodeMethod · 0.85
do_spec_decodeMethod · 0.85

Calls 7

step_graph_freeFunction · 0.85
draft_has_swa_layersFunction · 0.85
mask_align_upFunction · 0.85
dataMethod · 0.45
sizeMethod · 0.45

Tested by 3

mainFunction · 0.68
run_dflash_spec_decodeFunction · 0.68