diff --git a/CHANGELOG.md b/CHANGELOG.md index e063967..84678b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # Changelog +## Unreleased + +- Add optional `SPECULATION=tree` serving, with best-first DFlash proposals, + private BF16 parent states and explicit commit/discard ownership. The default + remains MTP; the switch is selected at startup. +- Preserve the existing prefill cache, scheduler, tool/SSE framing, reasoning + controls and cancellation. Tree lanes run cooperatively; linear chains reuse + the original graph verifier. +- Bound PLE hash history and add numerical, ownership and service tests. See + `docs/TREE-VERIFICATION.md` for measured scope and remaining acceptance gates. + ## 0.1.2 — 2026-10-28 (experimental) - Fix the GitHub Actions pip cache: point setup-python at `requirements-test.txt`. diff --git a/README.md b/README.md index 382359e..6df9baf 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,24 @@ stops the container if free RAM drops below 20 GiB or swap grows too fast. The server listens on port 8000 in the container and requires `Authorization: Bearer `. Endpoints: `/v1/models`, `/v1/chat/completions` (streaming and tools), `/v1/scheduler`. +### Optional tree speculation + +MTP remains the default. Select the experimental best-first tree backend at startup: + +```bash +SPECULATION=tree MODEL_DIR=... DFLASH_DIR=... KEY_FILE=... bash scripts/run_docker.sh +``` + +Use `SPECULATION=mtp` to select the original backend on the next start. Direct +entrypoint users can set `SPARK_SPECULATION=tree` or pass `--speculation tree` +with `SPARK_EXECUTION_MODE=continuous_mtp` and their usual model environment. +`/health` and `/v1/scheduler` report the selected mode. + +Tree mode retains streaming, tools, reasoning controls, paired prefill-cache +reuse and cancellation. Its private lanes run cooperatively, without GPU batch +verification; four-way full-context capacity has not been established. See +[implementation and validation](docs/TREE-VERIFICATION.md). + ## Layout ``` diff --git a/README.zh-CN.md b/README.zh-CN.md index f784d72..93ef26d 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -60,6 +60,23 @@ MODEL_DIR=... DFLASH_DIR=... KEY_FILE=... bash scripts/run_docker.sh 容器内监听 8000 端口，请求需带 `Authorization: Bearer `。接口：`/v1/models`、`/v1/chat/completions`（流式与工具）、`/v1/scheduler`。 +### 可选整树投机 + +默认继续使用 MTP。在启动时选择实验性的 best-first 整树后端： + +```bash +SPECULATION=tree MODEL_DIR=... DFLASH_DIR=... KEY_FILE=... bash scripts/run_docker.sh +``` + +下次启动设为 `SPECULATION=mtp` 可恢复原后端。直接运行入口时，可设置 +`SPARK_SPECULATION=tree` 或传入 `--speculation tree`，同时保留 +`SPARK_EXECUTION_MODE=continuous_mtp` 和原模型环境。`/health` 与 `/v1/scheduler` +会显示所选模式。 + +树模式保留流式、工具、思考控制、成对 prefill 缓存复用和取消流程。多请求使用私有 +状态合作式逐路执行，尚无 GPU 验证合批；四路满上下文容量未验收。 +详见[实现与验证记录](docs/TREE-VERIFICATION.md)。 + ## 测试 ```bash diff --git a/docs/TREE-VERIFICATION.md b/docs/TREE-VERIFICATION.md new file mode 100644 index 0000000..c0abbf5 --- /dev/null +++ b/docs/TREE-VERIFICATION.md @@ -0,0 +1,190 @@ +# Experimental tree verification + +The default serving runner uses its original MTP verifier. Select tree +speculation explicitly at startup with `SPECULATION=tree` in the Docker +launcher, or `SPARK_SPECULATION=tree` / `--speculation tree` in the direct +entrypoint. The direct entrypoint requires `continuous_mtp` mode and DFlash2. + +The scheduling idea was reviewed against [Uzu's source](https://github.com/trymirai/uzu/tree/f1231705cc7e26a205819da7c78e0110cfbac306). +The implementation uses Ember's existing CUDA arithmetic and current target +and DFlash2 checkpoints. It adds no model weights or draft training. + +## Two independent experimental paths + +`engine.verify_deferred(ids, state)` implements two-phase **linear** GDN +verification. It returns logits, hidden features and a transaction with explicit +`commit(consumed_input_count)` / `discard()` operations. Verification retains +convolved inputs and small convolution prefixes. Commit replays only the +accepted recurrent updates with the original decode kernel; projections and +experts are not repeated. + +For eight verification rows, the retained GDN snapshot payload drops from +13,074,432 to 656,896 bytes per layer. This accounting excludes model weights, +base state and output tensors. It is not total model memory usage. Full-model +paired verification-plus-commit timings showed parity, not a stable speed gain. + +`engine.verify_tree(tree, state, observer=None)` implements a small tree using +unique node rows. Shared ancestors are evaluated once; large projections, HC +and MoE use the unique rows together. Each GDN, QSA compressed-tail and PLE +n-gram/convolution history follows its parent. Positions follow ancestor depth. + +- `CandidateTree` permits one to eight nodes including the unconsumed anchor. +- The initial model backend permits at most four leaves. +- `draft.propose_tree(ctx, anchor)` reuses the DFlash2 candidate lattice with + two active branches and up to three draft levels (usually seven nodes). +- `draft.propose_tree(ctx, anchor, policy='best_first')`, or + `SPARK_TREE_POLICY=best_first`, spends up to eight nodes on the highest-scoring + available prefixes. It can extend seven draft levels while retaining at most + four leaves. The existing beam policy remains the default. These normalized + selector scores are proposal heuristics, not target probabilities. +- Selection follows target argmax. Proposal scores never accept a token. +- The predicting input is committed; EOS, correction and final bonus output + remain unconsumed. Output budgets are checked at every accepted node. +- `SPARK_TREE_GDN_KERNEL=1` enables three GDN launches per tree/layer: + convolution, recurrent updates and normalization. Parent caches are BF16. + Selected node-bank views are copied into independent committed storage. +- Without that flag, GDN calls the original single-node cores in parent order. + +PLE hashing reads only the required token-history tail (two tokens for the +current checkpoint), rather than copying the full history per candidate. +The parent convolution state retains its causal history. Tree depth is bounded +by the target checkpoint's remaining context; the final legal anchor can use +a root-only tree. + +The tree path retains per-node GDN caches and does **not** yet use the linear +path's deferred memory saving. Branched verification is eager, with no merged +tree CUDA Graph. Serving reuses the original linear graph verifier for an +unbranched chain. Private lanes run cooperatively on the existing model-owner +scheduler; no multi-lane tree GPU batch is implemented. + +## Numerical and ownership contracts + +The current checkpoint writes its GDN recurrence to BF16 between tokens; beta, +convolution products and core outputs have their own BF16 boundaries. An FP32 +closed-form or triangular solve is not an interchangeable implementation. +The tree kernel preserves the original recurrent expression/order and reads +each parent's materialized BF16 cache. + +Only one verification transaction may be pending on an engine. Another model +call discards the old lease. A committed path cannot cross siblings or publish +unaccepted state. Callers must serialize access and preserve the base state +until commit/discard. These interfaces implement greedy decoding only. + +`engine.verify_tree_reference(tree, state)` expands root-to-leaf paths into +existing private batch lanes and checks shared ancestor logits. It is a +correctness oracle and intentionally recomputes shared ancestors. + +## Validation tools + +- `tools/validate_gdn_deferred.py`: synthetic actual-shape GPU verification, + every accepted prefix, subsequent decode, immutable base and inference mode. +- `tools/validate_gdn_tree.py`: synthetic branched/linear topologies, every + node cache/core and subsequent decode against the original kernel. +- `tools/validate_linear_verification.py`: real current-checkpoint logits, + hidden, complete GDN/QSA/PLE state and generation checks. Tree-only runs use + `SPARK_VALIDATE_TREE_ONLY=1`; optional alternating generation timings use + `SPARK_TREE_TIMING=1`. +- `tools/tree_generation.py`: isolated full-generation harness, including + DFlash feature updates from **only the committed ancestor rows**. +- `tools/validate_tree_improvements.py`: current-weight serial-node/state/next-step + equality and alternating baseline/beam/best-first request and Agent timings. + +Run GPU tools with the reference dependencies and native libraries, exclusive +GPU ownership, a bounded memory cgroup and the memory guard. They do not open +an HTTP listener. Keep each attempt's source, exit status and results in a +separate evidence directory. + +Initial completed evidence (2026-10-29): 660 CPU tests passed, 13 tokenizer +tests skipped; the linear model passed all 22 accepted-prefix checks at 25 and +2,081 context tokens. The fused-GDN unique-node model passed all ten tree +checks, including a target path through the second branch. Synthetic fused +GDN passed all 14 topology/inference-mode cases, including subsequent decode. + +Four alternating-order full-request measurements per backend, after warm-up, +used the same current checkpoint and disabled the prefix cache. Every output +token matched. These are one-lane synchronous requests, including prefill and +draft, with no HTTP transport: + +| Request | Output tokens | Accepted runner | Eager tree | Tree GDN kernel | +| --- | ---: | ---: | ---: | ---: | +| Short Python function | 25 | 0.5532 s / 45.19 tok/s | 0.5737 s / 43.58 tok/s | 0.5556 s / 45.00 tok/s | +| Agent tool-dispatch module | 128 | 2.8510 s / 44.90 tok/s | 2.7602 s / 46.37 tok/s | 2.6593 s / 48.13 tok/s | + +The longer response improved about 7.2%; the short response was effectively +unchanged. One earlier cold smoke test took 1.04 s for the tree vs. 0.56 s for +the runner. Do not equate warmed results with cold-start latency or generalize +these two requests to Agent workflow throughput. The fused GDN microbenchmark +is a layer measurement, not an end-to-end speed ratio. + +`tools/validate_tree_agent.py` additionally reuses the existing virtual +read/patch/test and inspection fixtures with real normalization and tool +parsing, deterministic fixture call IDs, and identical prompt/output-token +hashes per turn. It remains a synchronous in-process test; HTTP/SSE and +concurrent scheduling require separate acceptance. + +These initial checks do not establish multi-step Agent throughput, model quality, concurrent +serving, cancellation/queue/cache behavior or four-lane 262K capacity. +`release_approved=false`. + +## Bounded-history and best-first update, 2026-10-30 + +688 CPU tests passed, with 13 tokenizer tests skipped and 1,189 subtests passed. +On the same current target and DFlash2 weights, all 36 GPU checks matched +serial decoding bit for bit: every node's logits/hidden, committed full state, +unchanged base state and subsequent decode. Cases include actual proposals, +an eight-node chain, asymmetric branches, four/six-node boundaries and output +budgets 1/2/8 at 25 and 2,081 context tokens. + +Four timed samples per cell after warm-up, alternating backend order, produced +identical output token IDs: + +| Single-lane synchronous request | Tokens | Existing runner | Beam tree | Best-first tree | +| --- | ---: | ---: | ---: | ---: | +| Short Python function | 25 | 45.21 tok/s | 44.85 tok/s | 58.72 tok/s | +| Tool-dispatch code module | 128 | 44.72 tok/s | 47.72 tok/s | 51.77 tok/s | +| Virtual read/patch/test/report workflow | 201 | 28.84 tok/s | 28.87 tok/s | 32.46 tok/s | +| Virtual file-inspection workflow | 297 | 33.09 tok/s | 33.24 tok/s | 34.24 tok/s | + +The Agent workflows passed 30/30 fixture runs (24 timed, six warm-up), with +identical prompt/output-token hashes at every turn. Their median workflow +times changed from 6.9699 to 6.1919 seconds and 8.9757 to 8.6734 seconds, +respectively. Throughput improved about 12.6% and 3.5% on these two fixtures. +These fixtures operate on in-memory files with simulated tests. They do not +measure real repository-repair success, HTTP/SSE, concurrent scheduling or +prefix-cache reuse. All timings include prefill and draft, with prefix caching +disabled; they are not pure decode TPS or a comparison with another engine. + +The code-module response needed 30 verification rounds instead of the beam +tree's 35. The short function needed four instead of six. This is one loaded +process and a small fixed workload, so the result is experimental. No service +or release was promoted, and `release_approved=false`. + +## Serving switch validation, 2026-10-30 + +The launcher defaults to `SPECULATION=mtp`; `SPECULATION=tree` selects +best-first tree decoding. Change modes on the next start, using a free container +name and the same model environment. Existing containers are retained by the +launcher. `/health`, `/v1/scheduler` and the startup receipt identify the mode; +request audits identify tree rounds and their node-consumption counts. + +Tree prefill retains the existing paired target/MTP checkpoint cache. Decode +keeps private target/DFlash state for each request and does not enter an MTP +fallback with stale draft state. EOS, output/reasoning limits and protocol +stops use the existing request control. Cancellation discards uncommitted +verification; close releases the per-request DFlash context. Selected graph +views are detached even with lazy cloning disabled. + +719 local CPU tests passed (13 tokenizer tests skipped). With the exact current +checkpoint, two separately started reference-image containers passed 14 MTP +and 16 tree HTTP checks: authentication, JSON/SSE, virtual tool workflows, +four concurrent private requests, repeated-prefix cache hits of 2,048 tokens, +reasoning-budget controls, TCP disconnect cancellation and recovery. Audits +confirmed actual tree decode and forced reasoning tokens. Both containers +stopped with exit 0; the GPU was released and the unrelated service was retained. + +These HTTP runs are functional smoke tests, not a new throughput study. The +earlier throughput tables use an in-process generation harness; HTTP uses a +different scheduling path and first requests can include graph capture. +Do not use those tables as the measured speedup of this service switch. +Full model quality, broad workloads, queue saturation and four full 262K +contexts remain unverified. `release_approved=false`. diff --git a/docs/VALIDATION.md b/docs/VALIDATION.md index d3ef355..b1a0e76 100644 --- a/docs/VALIDATION.md +++ b/docs/VALIDATION.md @@ -1,5 +1,8 @@ # Validation — 0.1.2 +The opt-in experimental branch has separate [tree-verification evidence](TREE-VERIFICATION.md). +It is not part of the v0.1.2 release results below. + Checks run on 2026-10-28 on a DGX Spark (SM121) and on macOS. | Check | Result | diff --git a/runtime/api_server.py b/runtime/api_server.py index 68d96e0..6432ae3 100644 --- a/runtime/api_server.py +++ b/runtime/api_server.py @@ -125,9 +125,13 @@ def __init__(self,options,queue_timeout=QUEUE_TIMEOUT): class Runtime: - def __init__(self,mode=None,max_active=None,max_queued=None,total_context=None,prepare_workers=None,mtp_batch=None): + def __init__(self,mode=None,max_active=None,max_queued=None,total_context=None,prepare_workers=None,mtp_batch=None,speculation=None): self.mode=mode or os.environ.get('SPARK_EXECUTION_MODE','serial_mtp') if self.mode not in ('serial_mtp','continuous_greedy','continuous_mtp'):raise ValueError('Unknown execution mode') + from speculation import settings + self.speculation=settings(speculation) + if self.speculation['mode']=='tree' and self.mode!='continuous_mtp': + raise ValueError('Tree speculation requires continuous_mtp mode') self.continuous=self.mode!='serial_mtp' self.mtp_enabled=self.mode!='continuous_greedy' batch_flag=os.environ.get('SPARK_MTP_BATCH','0') @@ -378,13 +382,18 @@ def load(self): self.batch_runner=GreedyBatchRunner(self.engine,self.prefix_cache) elif self.mode=='continuous_mtp': from engine.continuous_mtp import ContinuousMTPRunner - if self.mtp_batch: + tree_options={} + if self.speculation['mode']=='tree': + from engine.tree_runner import TreeMTPRunner + ContinuousMTPRunner=TreeMTPRunner + tree_options=dict(tree_policy=self.speculation['policy'],tree_kernel=self.speculation['gdn_kernel']) + elif self.mtp_batch: from engine.batched_mtp import BatchedMTPRunner ContinuousMTPRunner=BatchedMTPRunner from engine.speculative_policy import CostAwareDraftDepth self.batch_runner=ContinuousMTPRunner(self.engine,self.draft,self.prefix_cache, draft_tokens=DRAFT_TOKENS, - policy_factory=CostAwareDraftDepth if self.mtp_policy=='cost' else None) + policy_factory=CostAwareDraftDepth if self.mtp_policy=='cost' else None,**tree_options) if QSA_PREFILL_MODE not in ('packed','predecoded'):raise ValueError('Unknown SM121 QSA prefill mode.') for layer in self.engine.layers: if layer.layer_type=='full_attention':layer.attention.prefill_mode=QSA_PREFILL_MODE @@ -429,8 +438,9 @@ def load(self): queue_timeout_seconds=self.queue_timeout,max_pending_requests=self.max_active+self.max_queued, execution_mode=self.mode,max_active_requests=self.max_active,max_queued_requests=self.max_queued, mtp_enabled=self.mtp_enabled,owned_graph_cache=self.continuous, - scheduling='decode_priority_batch' if self.mtp_batch else 'cooperative_round_robin' if self.mode=='continuous_mtp' else self.mode, - batched_verification=self.mtp_batch,mtp_policy=self.mtp_policy if self.mtp_enabled else None, + scheduling='cooperative_round_robin' if self.speculation['mode']=='tree' else 'decode_priority_batch' if self.mtp_batch else 'cooperative_round_robin' if self.mode=='continuous_mtp' else self.mode, + batched_verification=self.mtp_batch and self.speculation['mode']!='tree',speculation=self.speculation, + mtp_policy=self.mtp_policy if self.mtp_enabled else None, batch_execution=self.batch_runner.snapshot() if self.mtp_batch else None, experimental=self.mode=='continuous_mtp', moe_strategy=strategy,actual_moe_strategy_counts=self.moe_strategy.observed if self.moe_strategy else {}, @@ -825,7 +835,8 @@ async def http_error(request,error): @app.get('/health') async def health(): - return JSONResponse(dict(status='ok' if runtime.ready else 'unavailable',model=MODEL),status_code=200 if runtime.ready else 503) + return JSONResponse(dict(status='ok' if runtime.ready else 'unavailable',model=MODEL, + speculation=getattr(runtime,'speculation',dict(mode='mtp'))),status_code=200 if runtime.ready else 503) @app.get('/v1/models') async def models(request:Request): @@ -842,7 +853,9 @@ async def scheduler_status(request:Request): authorize(request) scheduler=getattr(runtime,'scheduler',None) mode=getattr(runtime,'mode','serial_mtp') + speculation=getattr(runtime,'speculation',dict(mode='mtp')) return dict(execution_mode=mode, + speculation=speculation, context_window_per_request=getattr(runtime,'max_context',MAX_CONTEXT), max_output_tokens=getattr(runtime,'max_output',MAX_OUTPUT) or None, output_budget_policy=getattr(runtime,'output_budget_policy', @@ -852,8 +865,8 @@ async def scheduler_status(request:Request): prepare_workers=getattr(runtime,'prepare_workers',0), mtp_enabled=mode in ('serial_mtp','continuous_mtp'), reasoning_efforts=REASONING_EFFORTS if mode=='continuous_mtp' else {}, - scheduling='decode_priority_batch' if getattr(runtime,'mtp_batch',False) else 'cooperative_round_robin' if mode=='continuous_mtp' else mode, - batched_verification=getattr(runtime,'mtp_batch',False),experimental=mode=='continuous_mtp', + scheduling='cooperative_round_robin' if speculation['mode']=='tree' else 'decode_priority_batch' if getattr(runtime,'mtp_batch',False) else 'cooperative_round_robin' if mode=='continuous_mtp' else mode, + batched_verification=getattr(runtime,'mtp_batch',False) and speculation['mode']!='tree',experimental=mode=='continuous_mtp', **(scheduler.snapshot() if scheduler else dict(max_active=1, max_queued=getattr(runtime,'max_queued',1),events=[]))) diff --git a/runtime/engine/batch_mtp_model.py b/runtime/engine/batch_mtp_model.py index 57e789c..c86a181 100644 --- a/runtime/engine/batch_mtp_model.py +++ b/runtime/engine/batch_mtp_model.py @@ -229,6 +229,8 @@ def layer(self, index, pending, states, sections, positions=None, ids=None, cont def verify(self, input_ids, states, observer=None): """observer(layer_index, (hidden, block, injection)) sees every layer's output rows (all lanes, concatenated in lane order), as the single-lane engine's observer does.""" + discard=getattr(self.engine,'discard_verification',None) + if discard is not None:discard() lengths = [len(ids) for ids in input_ids]; sections = spans(lengths) if len(states) != len(sections) or len({id(s) for s in states}) != len(states): raise ValueError('Batch requires distinct private states') diff --git a/runtime/engine/candidate_tree.py b/runtime/engine/candidate_tree.py new file mode 100644 index 0000000..16ed269 --- /dev/null +++ b/runtime/engine/candidate_tree.py @@ -0,0 +1,142 @@ +"""Bounded candidate topology and greedy acceptance, independent of GPU state. + +Node zero is the already emitted, not-yet-consumed anchor. A row predicts the +next output token. Every emitted token commits its predicting input node, so +EOS, corrections and the final bonus token remain unconsumed exactly as in +the linear verifier. Proposal scores never decide target acceptance. +""" +from dataclasses import dataclass +import math +import heapq + + +@dataclass(frozen=True) +class TreeSelection: + consumed_nodes: tuple + emitted: tuple + reason: str + + +@dataclass(frozen=True) +class CandidateTree: + tokens: tuple + parents: tuple + + def __post_init__(self): + object.__setattr__(self,'tokens',tuple(self.tokens)) + object.__setattr__(self,'parents',tuple(self.parents)) + if not 1<=len(self.tokens)<=8 or len(self.parents)!=len(self.tokens): + raise ValueError('Expected 1..8 tree nodes including the anchor') + if any(type(p) is not int for p in self.parents) or self.parents[0]!=-1 or any(type(t) is not int or t<0 for t in self.tokens): + raise ValueError('Invalid anchor or token IDs') + edges=set() + for i in range(1,len(self.tokens)): + p=self.parents[i] + if type(p) is not int or not 0<=p=limit:return + values=list(candidates[depth]) + if not 0<=previous1:expand(0,0,0,0.) + while frontier and len(tokens)max_leaves:continue + node=len(tokens);tokens.append(token);parents.append(parent);children.append(False) + children[parent]=True;leaves=next_leaves;used.add(edge) + expand(node,depth+1,index,-negative) + return CandidateTree(tokens,parents) diff --git a/runtime/engine/continuous_mtp.py b/runtime/engine/continuous_mtp.py index c759b2b..4010476 100644 --- a/runtime/engine/continuous_mtp.py +++ b/runtime/engine/continuous_mtp.py @@ -261,6 +261,7 @@ def close(self): class ContinuousMTPRunner: + session_type=MTPSession def __init__(self, engine, draft, cache=None, *, draft_tokens='adaptive', policy_factory=None, low_sync=True): if draft_tokens not in (1, 2, 'adaptive'): @@ -287,7 +288,7 @@ def start(self, ids, budget, eos, chunk_size, on_tokens=None, on_progress=None, cache=True, should_cancel=None): policy = ((self.policy_factory or AdaptiveDraftDepth)() if self.draft_tokens == 'adaptive' else None) - session = MTPSession(ids, budget, eos, chunk_size, on_tokens, on_progress, + session = self.session_type(ids, budget, eos, chunk_size, on_tokens, on_progress, should_cancel, self.draft_tokens, policy, self.low_sync) session.prefill_memory = memory_summary(self.prefill_allocator.enabled,self.prefill_allocator.threshold_bytes) session.cache_enabled = cache diff --git a/runtime/engine/decoder.py b/runtime/engine/decoder.py index 9b426cc..ddbb009 100644 --- a/runtime/engine/decoder.py +++ b/runtime/engine/decoder.py @@ -40,7 +40,10 @@ def __init__(self,store,config,layer_index,activation_backend='torch',hc_backend config.get('norm_topk_prob',True),activation_backend=activation_backend,expert_backend=expert_backend,router_backend=router_backend,moe_finalize=moe_finalize) def forward(self,hidden_states,positions,prev_block_output=None,prev_injection=None, - input_ids_cpu=None,context_cpu=None,state=None,prefill=False,position_start=None,return_prefixes=False): + input_ids_cpu=None,context_cpu=None,state=None,prefill=False,position_start=None,return_prefixes=False, + defer_gdn=False): + if defer_gdn and (prefill or not return_prefixes): + raise ValueError('Deferred GDN requires decode verification') state=state or DecoderState() if (prev_block_output is None)!=(prev_injection is None): raise ValueError('Pending block output/injection must arrive together') @@ -61,7 +64,13 @@ def forward(self,hidden_states,positions,prev_block_output=None,prev_injection=N else: residual,mixed,injection=self.attn_hc.combine_and_mix(hidden_states,prev_block_output,prev_injection,return_prefixes) if self.layer_type=='linear_attention': - attention_result=self.attention(mixed,state.attention,prefill,return_prefixes) + if defer_gdn: + from .ops.gdn_deferred import verify_gdn + qkvz,ba=self.attention.project(mixed,preserve_rows=True) + tape=verify_gdn(self.attention,qkvz,ba,state.attention) + attention_result=(self.attention.out_linear(tape.output),tape,[tape]*len(mixed)) + else: + attention_result=self.attention(mixed,state.attention,prefill,return_prefixes) else: attention_result=self.attention(mixed,positions,state.attention,position_start=position_start,return_prefixes=return_prefixes,prefill_chunk=prefill) attention,attention_state=attention_result[:2] diff --git a/runtime/engine/dflash_draft.py b/runtime/engine/dflash_draft.py index 884c536..7bd586f 100644 --- a/runtime/engine/dflash_draft.py +++ b/runtime/engine/dflash_draft.py @@ -257,7 +257,7 @@ def extend(self, ctx, feats): ctx.length += n; ctx.written += n # ---- proposal ---- - def _propose_core(self, ctx, anchor, start): + def _propose_core(self, ctx, anchor, start, return_lattice=False): B, H, NH, KVH, D, W = self.B, self.H, self.NH, self.KVH, self.D, self.W qd, kd, win = NH*D, KVH*D, ctx.window ids = torch.cat((anchor, self.mask_ids)) @@ -298,6 +298,8 @@ def _propose_core(self, ctx, anchor, start): pred_ids = torch.cat((ids[:1, None].expand(1, self.TOPK), cand[:-1]), 0) preds = W['candidate_selector.predecessor_codebook'][pred_ids].float() # [s, K, r] score = vals_[:, None, :]+torch.einsum('lpr,lcr->lpc', preds*hp[:, None], keys_) + if return_lattice: + return cand,score ix = score[0, 0].argmax(-1); path = [ix] for e in range(1, B-1): # index_select keeps the walk on the device (a 0-d tensor subscript would sync to the host) @@ -323,3 +325,27 @@ def propose(self, ctx, anchor): return out ctx.prop.replay(); self.replays += 1 return ctx.prop_out.tolist() + + @torch.no_grad() + def propose_tree(self,ctx,anchor,node_budget=8,policy=None): + """Explicit reference proposal; eager lattice and CPU beam construction. + + Existing linear proposal graphs and weights are unchanged. This path + synchronizes to transfer the small candidate lattice, so it is not a + claim of a GPU-resident or accelerated tree proposer. + """ + from .candidate_tree import CandidateTree,tree_from_lattice,tree_from_lattice_best_first + policy=os.environ.get('SPARK_TREE_POLICY','beam') if policy is None else policy + if policy not in ('beam','best_first'):raise ValueError('Unknown tree proposal policy') + if type(node_budget) is not int or not 1<=node_budget<=8:raise ValueError('Invalid tree node budget') + depth=self.engine.config['max_position_embeddings']-ctx.length-1 + if depth<0:raise ValueError('Anchor exceeds target context limit') + self.flush(ctx) + if node_budget==1 or depth==0:return CandidateTree((int(anchor),),(-1,)) + ctx.anchor.fill_(int(anchor));ctx.start.fill_(ctx.length) + candidates,scores=self._propose_core(ctx,ctx.anchor,ctx.start,return_lattice=True) + if policy=='best_first': + return tree_from_lattice_best_first(int(anchor),candidates.tolist(),scores.tolist(), + node_budget=node_budget,max_depth=min(7,self.B-1,depth),max_leaves=4) + return tree_from_lattice(int(anchor),candidates.tolist(),scores.tolist(), + node_budget=node_budget,max_depth=min(3,self.B-1,depth),width=2) diff --git a/runtime/engine/model.py b/runtime/engine/model.py index 9090cdf..f33b5a7 100644 --- a/runtime/engine/model.py +++ b/runtime/engine/model.py @@ -100,6 +100,7 @@ def logits(self,hidden,chunk=4096): return output.float() def forward(self,input_ids,state=None,observer=None,return_hidden=False,return_prefixes=False,prefill_chunk=False,borrowed_graph_cache=False,owned_graph_cache=False,lazy_prefixes=False): + self.discard_verification() ids=[int(token) for token in input_ids] if not ids:raise ValueError('Expected at least one token') if any(token<0 or token>=self.config['vocab_size'] for token in ids):raise ValueError('Token is outside checkpoint vocabulary') @@ -200,6 +201,31 @@ def verify_tokens(self,input_ids,state,observer=None,borrowed_graph_cache=False, return self(input_ids,state,observer,return_hidden=True,return_prefixes=True,borrowed_graph_cache=borrowed_graph_cache, lazy_prefixes=lazy_prefixes) + def discard_verification(self): + pending=getattr(self,'_pending_verification',None) + if pending is not None: + pending.discard() + + def verify_deferred(self,input_ids,state,observer=None): + """Explicit experimental API; commit/discard before another model call. + + A subsequent model call invalidates any uncommitted transaction. Only + consumed input prefixes are committed; the last predicted token stays + unconsumed, as in verify_tokens. Existing serving paths are unchanged. + """ + from .verification import verify_linear + return verify_linear(self,input_ids,state,observer) + + def verify_tree_reference(self,tree,state): + """Experimental correctness oracle: expands paths, recomputes ancestors.""" + from .tree_verification import verify_tree_reference + return verify_tree_reference(self,tree,state) + + def verify_tree(self,tree,state,observer=None): + """Opt-in eager unique-node tree prototype; existing serving is unchanged.""" + from .tree_model import verify_tree + return verify_tree(self,tree,state,observer) + @torch.no_grad() def generate(self,prompt_ids,max_tokens=8): state=None; ids=prompt_ids; generated=[] diff --git a/runtime/engine/ops/gdn_commit.py b/runtime/engine/ops/gdn_commit.py new file mode 100644 index 0000000..b279bba --- /dev/null +++ b/runtime/engine/ops/gdn_commit.py @@ -0,0 +1,12 @@ +"""Commit independent GDN tapes through the bit-verified original kernel. + +The experimental cross-layer kernel failed strict BF16 equality and is not +part of the accepted path. All tapes are checked before any work is queued; +outputs remain private until the enclosing transaction publishes its state. +""" + + +def commit_many(tapes, count): + for tape in tapes: + tape.validate(count) + return [tape.commit(count) for tape in tapes] diff --git a/runtime/engine/ops/gdn_deferred.py b/runtime/engine/ops/gdn_deferred.py new file mode 100644 index 0000000..7de9917 --- /dev/null +++ b/runtime/engine/ops/gdn_deferred.py @@ -0,0 +1,110 @@ +"""Opt-in linear GDN verification without per-token recurrent snapshots. + +The existing kernels retain their BF16 products, beta, core and inter-token +state boundaries. Verification saves convolved inputs and small convolution +prefixes; commit replays only the accepted recurrent updates. No projections, +expert calls or model layers are rerun. This module does not enable itself. +""" +from dataclasses import dataclass, field + +import torch + +from .gdn import GDNState +from ..limits import VERIFY_WIDTHS + + +def _versions(state): + def version(tensor): + try: + return tensor._version + except RuntimeError: # inference_mode tensors have no version counter + return None + return tuple(version(x) for x in (state.conv, state.recurrent)) + + +@dataclass +class DeferredGDN: + output: torch.Tensor + model: object + initial: GDNState + conv_prefixes: torch.Tensor + convolved: torch.Tensor + ba: torch.Tensor + _status: str = field(default='open', init=False) + _versions: tuple = field(init=False, repr=False) + + def __post_init__(self): + self._versions = _versions(self.initial) + + @property + def rows(self): + return self.convolved.shape[0] + + @property + def retained_bytes(self): + # Model weights, base state and the returned output are excluded. + return sum(x.numel()*x.element_size() for x in + (self.conv_prefixes, self.convolved, self.ba)) + + def validate(self, count): + if self._status != 'open': + raise RuntimeError('Deferred GDN transaction is closed') + if type(count) is not int or not 1 <= count <= self.rows: + raise ValueError('Expected a nonempty accepted input prefix') + if self._versions != _versions(self.initial): + raise RuntimeError('Deferred GDN base state changed') + + def commit(self, count): + self.validate(count) + from .gdn_triton import _decode + recurrent = torch.empty_like(self.initial.recurrent) + # Keep output arithmetic live during replay to preserve the original + # kernel's reduction/layout contract. Its small result is discarded. + scratch = self.output.new_empty(count, 6144) + b, a = self.ba.split((48, 48), -1) + _decode[(4, 48)](self.convolved, a, b, self.model.A_log, + self.model.dt_bias, self.initial.recurrent, recurrent, scratch, + count, a.stride(0), b.stride(0), num_warps=1, num_stages=3) + state = GDNState(self.conv_prefixes[count-1].clone(), recurrent) + self._status = 'committed' + self._release() + return state + + def discard(self): + if self._status != 'open': + raise RuntimeError('Deferred GDN transaction is closed') + self._status = 'discarded' + self._release() + + def _release(self): + self.model = self.initial = self.conv_prefixes = self.convolved = self.ba = None + + +def verify_gdn(model, qkvz, ba, state): + """Inputs/base state are read-only until commit or discard on this stream.""" + if (not qkvz.is_cuda or qkvz.dtype != torch.bfloat16 or + qkvz.ndim != 2 or qkvz.shape[1] != 16384 or + qkvz.shape[0] not in VERIFY_WIDTHS or ba.shape != (len(qkvz), 96) or + ba.dtype != qkvz.dtype or ba.device != qkvz.device): + raise ValueError('Deferred GDN requires checkpoint CUDA BF16 verification rows') + if (state.conv.shape != (1, 10240, 3) or + state.recurrent.shape != (1, 48, 128, 128) or + any(x.dtype != qkvz.dtype or x.device != qkvz.device or not x.is_contiguous() + for x in (state.conv, state.recurrent))): + raise ValueError('Deferred GDN requires private checkpoint BF16 state') + from .gdn_triton import _conv, _decode, _gated_norm + rows = len(qkvz) + convolved = qkvz.new_empty(rows, 10240) + conv_prefixes = qkvz.new_empty(rows, *state.conv.shape) + _conv[(40,)](qkvz, state.conv, model.conv, convolved, conv_prefixes[-1], + qkvz.stride(0), 10240, rows, 256, True, conv_prefixes, True) + b, a = ba.split((48, 48), -1) + core = qkvz.new_empty(rows, 6144) + _decode[(4, 48)](convolved, a, b, model.A_log, model.dt_bias, + state.recurrent, None, core, rows, a.stride(0), b.stride(0), + STORE_FINAL=False, num_warps=1, num_stages=3) + normalized = torch.empty_like(core) + z = qkvz[:, 10240:] + _gated_norm[(rows*48,)](core, z, model.norm, normalized, z.stride(0), + model.eps, model.gate_activation == 'sigmoid', rows*48, 1, num_warps=1) + return DeferredGDN(normalized, model, state, conv_prefixes, convolved, ba) diff --git a/runtime/engine/ops/gdn_tree.py b/runtime/engine/ops/gdn_tree.py new file mode 100644 index 0000000..b0ba6ca --- /dev/null +++ b/runtime/engine/ops/gdn_tree.py @@ -0,0 +1,100 @@ +"""Small-tree GDN with the incumbent BF16 arithmetic and parent caches. + +One convolution, recurrence and normalization launch per layer. The recurrence +expression/order follows gdn_triton._decode; every parent cache is materialized +as BF16. This retains per-node states and is not an FP32 triangular solve. +""" +import torch +import triton as tr +import triton.language as tl + +from .gdn import GDNState +from .gdn_triton import _gated_norm + + +@tr.jit +def _conv_tree(X,S,W,P,Y,BANK,SX:tl.constexpr,T:tl.constexpr,B:tl.constexpr): + d=tl.program_id(0)*B+tl.arange(0,B) + w0=tl.load(W+d*4).to(tl.float32) + w1=tl.load(W+d*4+1).to(tl.float32) + w2=tl.load(W+d*4+2).to(tl.float32) + w3=tl.load(W+d*4+3).to(tl.float32) + for row in range(T): + parent=tl.load(P+row) + source=tl.where(parent<0,S+d*3,BANK+parent*10240*3+d*3) + first=tl.load(source).to(tl.float32) + second=tl.load(source+1).to(tl.float32) + third=tl.load(source+2).to(tl.float32) + current=tl.load(X+row*SX+d).to(tl.float32) + value=(first*w0).to(Y.dtype.element_ty).to(tl.float32) + value+=(second*w1).to(Y.dtype.element_ty).to(tl.float32) + value+=(third*w2).to(Y.dtype.element_ty).to(tl.float32) + value+=(current*w3).to(Y.dtype.element_ty).to(tl.float32) + activated=value/(1+tl.exp(-value)) + tl.store(Y+row*10240+d,activated) + tl.store(BANK+row*10240*3+d*3,second) + tl.store(BANK+row*10240*3+d*3+1,third) + tl.store(BANK+row*10240*3+d*3+2,current) + + +@tr.jit +def _decode_tree(X,A,B,ALOG,BIAS,S,P,BANK,O,T:tl.constexpr, + SA:tl.constexpr,SB:tl.constexpr): + head=tl.program_id(1); key_head=head//3 + kdim=tl.arange(0,128); vdim=tl.program_id(0)*32+tl.arange(0,32) + alog=tl.load(ALOG+head).to(tl.float32) + bias=tl.load(BIAS+head).to(tl.float32) + offset=head*16384+vdim[:,None]*128+kdim[None,:] + for row in range(T): + parent=tl.load(P+row) + source=tl.where(parent<0,S+offset,BANK+parent*48*16384+offset) + updated=tl.load(source).to(tl.float32) + q=tl.load(X+row*10240+key_head*128+kdim).to(tl.float32) + k=tl.load(X+row*10240+2048+key_head*128+kdim).to(tl.float32) + v=tl.load(X+row*10240+4096+head*128+vdim).to(tl.float32) + q=q/tl.sqrt(tl.sum(q*q)+1e-6) + k=k/tl.sqrt(tl.sum(k*k)+1e-6) + q=q*128**-.5 + a=tl.load(A+row*SA+head).to(tl.float32) + b=tl.load(B+row*SB+head).to(tl.float32) + x=a+bias + softplus=tl.where(x<=20,tl.log(1+tl.exp(x)),x) + g=-tl.exp(alog)*softplus + beta=tl.sigmoid(b).to(B.dtype.element_ty).to(tl.float32) + h=updated*tl.exp(g) + delta=(v-tl.sum(h*k[None,:],1))*beta + updated=h+delta[:,None]*k[None,:] + core=tl.sum(updated*q[None,:],1) + tl.store(O+row*6144+head*128+vdim,core) + tl.store(BANK+row*48*16384+offset,updated) + + +def run_tree(model,qkvz,ba,state,tree,parents=None): + from ..candidate_tree import CandidateTree + if (not isinstance(tree,CandidateTree) or not qkvz.is_cuda or + qkvz.shape!=(len(tree.tokens),16384) or ba.shape!=(len(tree.tokens),96) or + any(x.dtype!=torch.bfloat16 or x.device!=qkvz.device for x in + (qkvz,ba,state.conv,state.recurrent)) or not qkvz.is_contiguous() or + state.conv.shape!=(1,10240,3) or state.recurrent.shape!=(1,48,128,128) or + not state.conv.is_contiguous() or not state.recurrent.is_contiguous()): + raise ValueError('Expected current-checkpoint CUDA BF16 tree inputs and private state') + if parents is None:parents=torch.tensor(tree.parents,device=qkvz.device,dtype=torch.int32) + if (parents.shape!=(len(tree.tokens),) or parents.device!=qkvz.device or + parents.dtype!=torch.int32 or not parents.is_contiguous()): + raise ValueError('Expected tree parent indices') + rows=len(tree.tokens) + convolved=qkvz.new_empty(rows,10240) + conv=qkvz.new_empty(rows,*state.conv.shape) + recurrent=qkvz.new_empty(rows,*state.recurrent.shape) + core=qkvz.new_empty(rows,6144);output=torch.empty_like(core) + _conv_tree[(40,)](qkvz,state.conv,model.conv,parents,convolved,conv,qkvz.stride(0),rows,256) + b,a=ba.split((48,48),-1) + _decode_tree[(4,48)](convolved,a,b,model.A_log,model.dt_bias,state.recurrent,parents, + recurrent,core,rows,a.stride(0),b.stride(0),num_warps=1,num_stages=3) + z=qkvz[:,10240:] + _gated_norm[(rows*48,)](core,z,model.norm,output,z.stride(0),model.eps, + model.gate_activation=='sigmoid',rows*48,1,num_warps=1) + states=[] + for node in range(rows): + value=GDNState(conv[node],recurrent[node]);value.tree_view=True;states.append(value) + return output,states diff --git a/runtime/engine/ops/gdn_triton.py b/runtime/engine/ops/gdn_triton.py index 37f9dee..64f61ea 100644 --- a/runtime/engine/ops/gdn_triton.py +++ b/runtime/engine/ops/gdn_triton.py @@ -138,7 +138,7 @@ def _gated_norm(X,Z,W,Y,SZ:tl.constexpr,EPS:tl.constexpr,SIGMOID:tl.constexpr, @tr.jit def _decode(X,A,B,ALOG,BIAS,S,UPDATED,O,T:tl.constexpr=1, SA:tl.constexpr=48,SB:tl.constexpr=48,PREFIX=None, - CAPTURE_PREFIX:tl.constexpr=False): + CAPTURE_PREFIX:tl.constexpr=False,STORE_FINAL:tl.constexpr=True): head=tl.program_id(1); key_head=head//3 kdim=tl.arange(0,128); vdim=tl.program_id(0)*32+tl.arange(0,32) alog=tl.load(ALOG+head).to(tl.float32) @@ -169,7 +169,8 @@ def _decode(X,A,B,ALOG,BIAS,S,UPDATED,O,T:tl.constexpr=1, # tokens. The next row must observe exactly that materialization. if row+14:raise ValueError('Prototype supports at most four leaves') + if (not state.tokens or len(state.layers)!=len(engine.layers) or + any(t>=engine.config['vocab_size'] for t in tree.tokens)): + raise ValueError('Invalid checkpoint tree or base state') + start=len(state.tokens) + if max(tree.positions(start))>=engine.config['max_position_embeddings']: + raise ValueError('Tree depth exceeds context limit') + if (getattr(engine,'tree_chain_fastpath',False) and len(tree.tokens)>1 and + tree.parents==(-1,)+tuple(range(len(tree.tokens)-1))): + from .tree_verification import verify_tree_reference + tx=verify_tree_reference(engine,tree,state,observer) + tx.implementation='linear_chain_graphs' + return tx + engine.discard_verification() + device=engine.embedding.device + tokens=torch.tensor(tree.tokens,device=device,dtype=torch.long) + positions=torch.tensor(tree.positions(start),device=device,dtype=torch.long) + kernel=getattr(engine,'tree_gdn_kernel',os.environ.get('SPARK_TREE_GDN_KERNEL')=='1') + parents=torch.tensor(tree.parents,device=device,dtype=torch.int32) if kernel else None + pending=(engine.embedding[tokens].repeat(1,engine.config['hc_count']),None,None) + node_layers=[[] for _ in tree.tokens] + for index,(layer,base) in enumerate(zip(engine.layers,state.layers)): + hidden,block,injection=pending + if layer.ple is not None: + if block is not None: + hidden=layer.attn_hc.combine(hidden,block,injection) + block=injection=None + enhancement,ple=ple_nodes(layer,hidden,tree,state.tokens,base.ple) + hidden=hidden+enhancement + else:ple=[base.ple]*len(tree.tokens) + if block is None: + residual,mixed,injection=layer.attn_hc.mix(hidden,True) + else: + residual,mixed,injection=layer.attn_hc.combine_and_mix(hidden,block,injection,True) + attention,states=attention_nodes(layer,mixed,positions,tree,base.attention,start,parents) + residual,mixed,injection=layer.mlp_hc.combine_and_mix(residual,attention,injection,True) + pending=(residual,layer.moe(mixed,preserve_rows=True),injection) + if observer is not None:observer(index,pending) + for node,(att,pl) in enumerate(zip(states,ple)): + node_layers[node].append(DecoderState(att,pl)) + multi,sampled,_=engine.final_mixer.combine_and_mix(*pending,True) + if observer is not None:observer('final',(multi,sampled,None)) + logits=engine.logits(sampled) + nodes=[EngineState(layers,list(state.tokens)+[tree.tokens[n] for n in tree.path(node)]) + for node,layers in enumerate(node_layers)] + outputs=[];paths=tree.paths() + for path in paths: + indices=torch.tensor(path,device=device,dtype=torch.long) + final=nodes[path[-1]] + branch=EngineState(final.layers,final.tokens,[nodes[n] for n in path]) + outputs.append((logits.index_select(0,indices),branch, + ForwardHidden(multi.index_select(0,indices),sampled.index_select(0,indices)))) + tx=TreeVerification(engine,tree,state,paths,outputs) + tx.implementation=('unique_nodes_tree_gdn_kernel' if parents is not None + else 'unique_nodes_parent_states_eager') + engine._pending_verification=tx + return tx diff --git a/runtime/engine/tree_runner.py b/runtime/engine/tree_runner.py new file mode 100644 index 0000000..f1aea6c --- /dev/null +++ b/runtime/engine/tree_runner.py @@ -0,0 +1,140 @@ +"""Opt-in tree sessions on the existing owner-thread scheduler and cache. + +Prefill keeps paired MTP cache checkpoints. Decode uses private target/DFlash +states only; lanes execute cooperatively and never mix their tree histories. +""" +import os +import time + +import torch + +from .candidate_tree import CandidateTree +from .continuous_mtp import ContinuousMTPRunner,MTPSession +from .dflash_draft import DFlashDraft,LAYERS +from .prefix_cache import clone_state + + +class TreeSession(MTPSession): + def result(self,cache_bytes): + result=super().result(cache_bytes) + result['draft_tokens']=None + result['speculative'].update(strategy='tree',scheduling='cooperative_round_robin', + batched_verification=False,policy=dict(strategy=self.tree_policy,max_nodes=8,max_leaves=4), + tree_consumption_rate=self.stats.get('tree_consumed_nodes',0)/self.stats['tree_nodes'] + if self.stats.get('tree_nodes') else None) + return result + + def close(self): + self.dflash_ctx=None + super().close() + + +class TreeMTPRunner(ContinuousMTPRunner): + session_type=TreeSession + + def __init__(self,*args,dflash=None,tree_policy='best_first',tree_kernel=True,**kwargs): + if tree_policy not in ('beam','best_first'):raise ValueError('Unknown tree proposal policy') + if type(tree_kernel) is not bool:raise ValueError('Expected a boolean tree kernel option') + super().__init__(*args,**kwargs) + self.tree_policy=tree_policy;self.engine.tree_gdn_kernel=tree_kernel + self.engine.tree_chain_fastpath=True + if dflash is None: + directory=os.environ.get('SPARK_DFLASH_DIR') + if not directory:raise ValueError('Tree speculation requires DFlash2') + head=getattr(self.draft,'head',None) if os.environ.get('SPARK_DFLASH_COMPACT_HEAD','1')=='1' else None + dflash=DFlashDraft(directory,self.engine,int(os.environ.get('SPARK_DFLASH_BLOCK','8')),head=head) + self.dflash=dflash + + def snapshot(self): + return dict(enabled=True,strategy='tree',policy=self.tree_policy,max_nodes=8,max_leaves=4, + scheduling='cooperative_round_robin',batched_verification=False,mixed_prefill_decode=False) + + def start(self,*args,**kwargs): + session=super().start(*args,**kwargs);session.tree_policy=self.tree_policy + session.stats.update(tree_rounds=0,tree_nodes=0,tree_consumed_nodes=0) + return session + + def prefill(self,session): + super().prefill(session) + if session.phase!='prefill': + # Cached prefill snapshots own their paired MTP states. This + # request never reenters MTP decode, so no stale fallback exists. + session.draft_state=session.draft_multi=session.proposal=None + + def _decode_one(self,session,emission_limit=None): + self._step(session,emission_limit) + + def force_next(self,session,token): + error=self.decode_batch([session],forced_tokens=[token])[0] + if error is not None:raise error + + def _step(self,session,emission_limit=None,forced=None): + self._cancel(session);self._wait();begin=time.perf_counter() + remaining=session.budget-len(session.generated) + if emission_limit is not None:remaining=min(remaining,emission_limit) + ctx=getattr(session,'dflash_ctx',None) + if ctx is None or ctx.length!=len(session.state.tokens):raise RuntimeError('Tree draft history is misaligned') + taps={} + def observe(index,value): + if isinstance(index,int) and index in LAYERS:taps[index]=self.dflash.tap(*value) + known=session.generated[-1];tree=None + if forced is not None: + # Ordinary forward can mutate its input container. The forced + # branch stays disposable until the cancellation check below. + _,state=self.engine([known],clone_state(session.state),observer=observe,owned_graph_cache=True) + nodes=(0,);emitted=(forced,);reason='forced' + else: + tree=(CandidateTree((known,),(-1,)) if remaining==1 else + self.dflash.propose_tree(ctx,known,node_budget=8,policy=self.tree_policy)) + tx=self.engine.verify_tree(tree,session.state,observe) + try: + selected=tx.select(remaining,session.eos) + self._cancel(session) + state=tx.commit(selected) + nodes=selected.consumed_nodes;emitted=selected.emitted;reason=selected.reason + finally: + if tx.status=='open':tx.discard() + indices=torch.tensor(nodes,device=self.engine.embedding.device,dtype=torch.long) + features=torch.cat([taps[i].index_select(0,indices) for i in LAYERS],-1) + self._cancel(session) + self.dflash.push(ctx,features);session.accept(emitted,state) + if ctx.length!=len(state.tokens):raise RuntimeError('Committed tree draft history is misaligned') + session.stats['tree_rounds']+=1 + session.stats['tree_nodes']+=len(tree.tokens) if tree is not None else 1 + session.stats['tree_consumed_nodes']+=len(nodes) + if forced is not None:session.stats['forced_control_tokens']+=1 + elif len(tree.tokens)==1:session.stats['single_steps']+=1 + else: + edges=set(zip(tree.parents[1:],tree.tokens[1:])) + accepted=sum((node,token) in edges for node,token in zip(nodes,emitted)) + corrections=int(reason=='miss');proposed=len(tree.tokens)-1 + session.stats['proposed']+=proposed;session.stats['accepted']+=accepted + session.stats['corrections']+=corrections;session.stats['discarded']+=proposed-accepted-corrections + session.stats['bonus_tokens']+=int(reason=='leaf') + self._wait() + session.iterations.append(dict(mode='reasoning_control' if forced is not None else 'tree', + emitted=len(emitted),verification_rows=len(tree.tokens) if tree is not None else 1, + tree_depth=max(map(len,tree.paths()))-1 if tree is not None else 0, + tree_reason=reason,batch_size=1,seconds=time.perf_counter()-begin)) + self._cancel(session) + + @torch.no_grad() + def decode_batch(self,sessions,*,emission_limits=None,forced_tokens=None): + count=len(sessions) + if (not 1<=count<=4 or len({id(s) for s in sessions})!=count or + any(s.closed or s.phase!='decode' for s in sessions)): + raise ValueError('Expected one to four distinct decode-ready tree sessions') + limits=[None]*count if emission_limits is None else list(emission_limits) + forced=[None]*count if forced_tokens is None else list(forced_tokens) + if len(limits)!=count or len(forced)!=count:raise ValueError('Control lane counts differ') + if any(v is not None and (type(v) is not int or v<1) for v in limits):raise ValueError('Expected positive emission limits') + if any(t is not None and (type(t) is not int or t<0 or t in s.eos) for s,t in zip(sessions,forced)): + raise ValueError('Expected non-EOS control tokens') + errors=[None]*count + for index,session in enumerate(sessions): + try:self._step(session,limits[index],forced[index]) + except Exception as error: + errors[index]=error + if isinstance(error,RuntimeError) and any(s in str(error).lower() for s in ('cuda','nonfinite')): + return [e or error for e in errors] + return errors diff --git a/runtime/engine/tree_verification.py b/runtime/engine/tree_verification.py new file mode 100644 index 0000000..6418cf8 --- /dev/null +++ b/runtime/engine/tree_verification.py @@ -0,0 +1,98 @@ +"""Correctness reference for bounded trees through existing private lanes. + +Root-to-leaf paths are expanded into independent batch lanes. Shared ancestors +are intentionally recomputed and checked for identical logits. This establishes +tree semantics with the existing QSA/PLE implementations; it is NOT the merged +node verifier and makes no tree-speed claim. +""" +import torch + +from .batch_mtp_model import BatchMTPModel, materialize_views +from .candidate_tree import CandidateTree, TreeSelection +from .model import EngineState, ForwardHidden + + +class TreeVerification: + def __init__(self,engine,tree,base,paths,outputs): + self.engine=engine;self.tree=tree;self.base=base;self.base_tokens=tuple(base.tokens) + self.paths=paths;self.outputs=outputs;self.status='open' + logits={};hidden={};sampled={} + for path,(rows,_,features) in zip(paths,outputs): + if not bool(torch.isfinite(rows).all()):raise RuntimeError('Nonfinite tree logits') + for row,node in enumerate(path): + if node in logits and not torch.equal(rows[row],logits[node]): + raise RuntimeError('Shared ancestor logits differ across reference lanes') + logits[node]=rows[row];hidden[node]=features.multi[row] + sampled[node]=features.sampled[row] + self.logits=torch.stack([logits[n] for n in range(len(tree.tokens))]) + self.hidden=ForwardHidden(torch.stack([hidden[n] for n in range(len(tree.tokens))]), + torch.stack([sampled[n] for n in range(len(tree.tokens))])) + self.predictions=tuple(self.logits.argmax(-1).tolist()) + + def select(self,budget,eos=()): + return self.tree.accept_argmax(self.predictions,budget,eos) + + def commit(self,selection): + if self.status!='open' or getattr(self.engine,'_pending_verification',None) is not self: + raise RuntimeError('Tree verification is closed or superseded') + if not isinstance(selection,TreeSelection):raise ValueError('Expected TreeSelection') + nodes=tuple(selection.consumed_nodes) + if (not nodes or nodes[0]!=0 or any(type(n) is not int or not 0<=n4:raise ValueError('Reference verifier supports at most four leaves') + if (not state.tokens or len(state.layers)!=len(engine.layers) or + any(t>=engine.config['vocab_size'] for t in tree.tokens)): + raise ValueError('Invalid checkpoint tree or base state') + if max(tree.positions(len(state.tokens)))>=engine.config['max_position_embeddings']: + raise ValueError('Tree depth exceeds context limit') + engine.discard_verification() + ids=[[tree.tokens[n] for n in path] for path in paths] + if len(paths)==1: + outputs=[engine.verify_tokens(ids[0],state,lazy_prefixes=True, + **({'observer':observer} if observer is not None else {}))] + else: + if observer is not None:raise ValueError('Reference observers require one linear path') + if not hasattr(engine,'_tree_reference_model'): + engine._tree_reference_model=BatchMTPModel(engine) + states=[EngineState(list(state.layers),list(state.tokens)) for _ in paths] + outputs=engine._tree_reference_model.verify(ids,states) + tx=TreeVerification(engine,tree,state,paths,outputs) + engine._pending_verification=tx + return tx diff --git a/runtime/engine/verification.py b/runtime/engine/verification.py new file mode 100644 index 0000000..1a7150e --- /dev/null +++ b/runtime/engine/verification.py @@ -0,0 +1,151 @@ +"""Explicit two-phase linear verification; never activated by import. + +Large GDN recurrence banks are replaced by a short-lived verification tape. +QSA and PLE keep the existing immutable per-prefix states. CUDA Graph tapes +are borrowed only until commit/discard or the next model call; no borrowed +tensor is published as a committed GDN state. +""" +import numpy as np +import torch + +from .decoder import DecoderState +from .decode_graph import DecodeGraphBlock +from .limits import VERIFY_WIDTHS +from .model import EngineState, ForwardHidden +from .ops.gdn_deferred import DeferredGDN + + +class DeferredGraphBlock(DecodeGraphBlock): + def _run(self): + value=self.pending + self.tapes=[];self.intermediates=[] + for layer,state in zip(self.layers,self.states): + value,updated=layer(value[0],None,value[1],value[2],state=state, + prefill=False,return_prefixes=True,defer_gdn=True) + self.tapes.append(updated.attention) + self.intermediates.append(value) + return value + + def forward(self,pending,states): + if self.graph is None:self._capture(pending,states) + else:self._copy_inputs(pending,states) + self.graph.replay();self.replays+=1 + updated=[] + for source,captured in zip(states,self.tapes): + tape=DeferredGDN(captured.output,captured.model,source.attention, + captured.conv_prefixes,captured.convolved,captured.ba) + updated.append(DecoderState(tape,prefixes=[DecoderState(tape) for _ in range(tape.rows)])) + return self.output,updated + + def __call__(self,pending,states): + return self.forward(pending,states) + + +class LinearVerification: + def __init__(self,engine,ids,base): + self.engine=engine;self.ids=tuple(ids);self.base=base + self.base_tokens=tuple(base.tokens);self.layers=[] + self.logits=self.hidden=None;self.status='open' + + @property + def retained_gdn_bytes(self): + return sum(layer.attention.retained_bytes for layer in self.layers + if isinstance(layer.attention,DeferredGDN)) + + def _validate(self,count): + if self.status!='open' or getattr(self.engine,'_pending_verification',None) is not self: + raise RuntimeError('Verification is closed or superseded') + if type(count) is not int or not 1<=count<=len(self.ids): + raise ValueError('Expected a nonempty consumed input prefix') + if tuple(self.base.tokens)!=self.base_tokens: + raise RuntimeError('Verification base token history changed') + for layer in self.layers: + if isinstance(layer.attention,DeferredGDN):layer.attention.validate(count) + + def commit(self,count): + self._validate(count) + try: + tapes=[layer.attention for layer in self.layers if isinstance(layer.attention,DeferredGDN)] + if tapes: + from .ops.gdn_commit import commit_many + committed=iter(commit_many(tapes,count)) + else:committed=iter(()) + layers=[] + for layer in self.layers: + prefix=layer.prefixes[count-1] + attention=(next(committed) if isinstance(layer.attention,DeferredGDN) + else prefix.attention) + layers.append(DecoderState(attention,prefix.ple)) + result=EngineState(layers,list(self.base_tokens)+list(self.ids[:count])) + self.status='committed' + return result + except BaseException: + self.status='failed' + raise + finally: + self._release() + + def discard(self): + if self.status!='open':raise RuntimeError('Verification is already closed') + self.status='discarded';self._release() + + def _release(self): + for layer in self.layers: + if isinstance(layer.attention,DeferredGDN) and layer.attention._status=='open': + layer.attention.discard() + if getattr(self.engine,'_pending_verification',None) is self: + self.engine._pending_verification=None + self.layers=[];self.base=None + + +@torch.no_grad() +def verify_linear(engine,input_ids,state,observer=None): + ids=list(input_ids) + if (len(ids) not in VERIFY_WIDTHS or not state.tokens or + len(state.layers)!=len(engine.layers) or + any(type(t) is not int or not 0<=tengine.config['max_position_embeddings']: + raise ValueError('Verification exceeds context limit') + engine.discard_verification() + tx=LinearVerification(engine,ids,state);engine._pending_verification=tx + try: + device=engine.embedding.device + token_ids=torch.tensor(ids,dtype=torch.long,device=device) + positions=torch.arange(start,start+len(ids),device=device) + preceding=engine.config['ngram_size']-1 + context=([engine.config['eos_token_id']]*preceding+state.tokens)[-preceding:] if preceding else [] + context=np.asarray([context],dtype=np.int64);ids_cpu=np.asarray(ids,dtype=np.int64) + pending=(engine.embedding[token_ids].repeat(1,engine.config['hc_count']),None,None) + index=0 + while index