vision: images ride the DSpark speculative loop
Browse files
recipe/base_evaluator.patch
CHANGED
|
@@ -1,15 +1,16 @@
|
|
| 1 |
--- /tmp/base_eval.bak 2026-07-03 13:59:36.997249952 -0400
|
| 2 |
-
+++ deepspec/eval/base_evaluator.py 2026-07-03
|
| 3 |
-
@@ -317,6 +317,
|
| 4 |
propose: Callable[..., DraftProposal],
|
| 5 |
update: Callable[[Any, VerificationResult], None],
|
| 6 |
post_verify: Callable[[DraftProposal, VerificationResult], None] | None = None,
|
| 7 |
+ prefill_keep_hidden_layers: list[int] | None = None,
|
| 8 |
+ stream_callback=None,
|
|
|
|
| 9 |
) -> SimpleNamespace:
|
| 10 |
"""Speculative-decoding loop.
|
| 11 |
|
| 12 |
-
@@ -343,14 +
|
| 13 |
from deepspec.eval.windowed_cache import build_target_cache
|
| 14 |
past_key_values_target = build_target_cache(target_model, pad=max(64, int(max_proposal_tokens)+8))
|
| 15 |
|
|
@@ -22,7 +23,7 @@
|
|
| 22 |
- logits_to_keep=1,
|
| 23 |
- )
|
| 24 |
+ _chunk = int(os.environ.get("DSPARK_PREFILL_CHUNK", "4096"))
|
| 25 |
-
+ if num_input_tokens > _chunk and prefill_keep_hidden_layers is not None:
|
| 26 |
+ # chunked prefill: bound activation memory (never materialize all layers x all positions);
|
| 27 |
+ # keep only the draft target layers hidden states, concatenated across chunks.
|
| 28 |
+ _keep = sorted(set(int(l) for l in prefill_keep_hidden_layers))
|
|
@@ -62,11 +63,12 @@
|
|
| 62 |
+ use_cache=True,
|
| 63 |
+ output_hidden_states=True,
|
| 64 |
+ logits_to_keep=1,
|
|
|
|
| 65 |
+ )
|
| 66 |
|
| 67 |
output_ids[:, :num_input_tokens] = input_ids
|
| 68 |
output_ids[:, num_input_tokens : num_input_tokens + 1] = sample_from_probs(
|
| 69 |
-
@@ -384,6 +
|
| 70 |
)
|
| 71 |
|
| 72 |
while start < max_length:
|
|
@@ -75,7 +77,7 @@
|
|
| 75 |
proposal = propose(
|
| 76 |
context=context,
|
| 77 |
output_ids=output_ids,
|
| 78 |
-
@@ -429,6 +
|
| 79 |
if has_stop_token(new_token_ids, stop_token_ids):
|
| 80 |
break
|
| 81 |
|
|
|
|
| 1 |
--- /tmp/base_eval.bak 2026-07-03 13:59:36.997249952 -0400
|
| 2 |
+
+++ deepspec/eval/base_evaluator.py 2026-07-03 20:44:32.070666055 -0400
|
| 3 |
+
@@ -317,6 +317,9 @@
|
| 4 |
propose: Callable[..., DraftProposal],
|
| 5 |
update: Callable[[Any, VerificationResult], None],
|
| 6 |
post_verify: Callable[[DraftProposal, VerificationResult], None] | None = None,
|
| 7 |
+ prefill_keep_hidden_layers: list[int] | None = None,
|
| 8 |
+ stream_callback=None,
|
| 9 |
+
+ prefill_mm: dict | None = None,
|
| 10 |
) -> SimpleNamespace:
|
| 11 |
"""Speculative-decoding loop.
|
| 12 |
|
| 13 |
+
@@ -343,14 +346,49 @@
|
| 14 |
from deepspec.eval.windowed_cache import build_target_cache
|
| 15 |
past_key_values_target = build_target_cache(target_model, pad=max(64, int(max_proposal_tokens)+8))
|
| 16 |
|
|
|
|
| 23 |
- logits_to_keep=1,
|
| 24 |
- )
|
| 25 |
+ _chunk = int(os.environ.get("DSPARK_PREFILL_CHUNK", "4096"))
|
| 26 |
+
+ if num_input_tokens > _chunk and prefill_keep_hidden_layers is not None and not prefill_mm:
|
| 27 |
+ # chunked prefill: bound activation memory (never materialize all layers x all positions);
|
| 28 |
+ # keep only the draft target layers hidden states, concatenated across chunks.
|
| 29 |
+ _keep = sorted(set(int(l) for l in prefill_keep_hidden_layers))
|
|
|
|
| 63 |
+ use_cache=True,
|
| 64 |
+ output_hidden_states=True,
|
| 65 |
+ logits_to_keep=1,
|
| 66 |
+
+ **(prefill_mm or {}),
|
| 67 |
+ )
|
| 68 |
|
| 69 |
output_ids[:, :num_input_tokens] = input_ids
|
| 70 |
output_ids[:, num_input_tokens : num_input_tokens + 1] = sample_from_probs(
|
| 71 |
+
@@ -384,6 +422,8 @@
|
| 72 |
)
|
| 73 |
|
| 74 |
while start < max_length:
|
|
|
|
| 77 |
proposal = propose(
|
| 78 |
context=context,
|
| 79 |
output_ids=output_ids,
|
| 80 |
+
@@ -429,6 +469,8 @@
|
| 81 |
if has_stop_token(new_token_ids, stop_token_ids):
|
| 82 |
break
|
| 83 |
|