| """Validate the native single-request MTP state commit without changing it.""" |
| import json |
| from pathlib import Path |
| import torch |
|
|
| def install(config): |
| from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2 |
| if getattr(EAGLEWorkerV2,'_strict_commit_audit',False):return |
| EAGLEWorkerV2._strict_commit_audit=True |
| original=EAGLEWorkerV2._mamba_verify_update |
| path=Path(config['output']+'.commit') |
| def checked(self,batch,lens,indices,bs): |
| result=original(self,batch,lens,indices,bs) |
| if bs==1 and not batch.forward_mode.is_idle(): |
| backend=self.target_worker.model_runner.attn_backend.linear_attn_backend |
| pool=backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() |
| slot=int(backend.forward_metadata.mamba_cache_indices[0]);step=int(lens[0])-1 |
| ok=torch.equal(pool.temporal[:,slot],pool.intermediate_ssm[:,0,step]) and all(torch.equal(live[:,slot],buf[:,0,step]) for live,buf in zip(pool.conv,pool.intermediate_conv_window)) |
| if not ok:raise RuntimeError('Strict MTP state commit mismatch') |
| with path.open('a') as f:f.write(json.dumps({'accepted_nodes':step+1,'state_equal':ok})+'\n') |
| return result |
| EAGLEWorkerV2._mamba_verify_update=checked |
|
|
