diff --git a/TECHNICAL_REPORTS/2269-decode-readback-sampler-recovery-20261011.en.md b/TECHNICAL_REPORTS/2269-decode-readback-sampler-recovery-20261011.en.md new file mode 100644 index 000000000..eed426031 --- /dev/null +++ b/TECHNICAL_REPORTS/2269-decode-readback-sampler-recovery-20261011.en.md @@ -0,0 +1,95 @@ +# Technical Report: PR #2269 — Contain Decode Readback Failures and Preserve Sampler State + +**Date**: 2026-10-11 + +**Status**: Implementation and required local validation completed. CI remains pending at report preparation. + +**Languages**: Rust, C++ + +**Risk Level**: High. Changes cross the lazy GPU graph, CXX error boundary, speculative KV teardown, and mutable sampler/RNG state. + +## Executive Summary + +Decode host reads now wait through fallible CXX entrypoints so backend evaluation errors can fail the affected request instead of escaping an infallible bridge. Failed lookahead submissions commit their already-drawn pending token before synchronous fallback, while failed row draws restore mirostat and adaptive-p feedback. + +Seeded fused parity additionally requires restoring the RNG key consumed while constructing the failed next submission. A snapshot of the pinned MLX thread-local key sequence supplies that rollback without adding a GPU evaluation or scheduling call. Focused fault tests, the broader suites and the newly strengthened stopping-boundary mutation checks passed; the restored implementation passed all 13 fault regressions again. + +## 1. Problem Statement + +An asynchronous backend error may first surface when the host waits for a token, selected probability, or sorted sampler distribution. Several decode paths reached non-`Result` CXX reads before a fallible wait, allowing a C++ exception to terminate the process. Guarding only the final synchronous token read did not protect mirostat or adaptive-p, which read earlier inside sampling. + +A second fault was state advancement after speculative failure. Dropping a pending token and sampling it again could consume another random key or advance feedback twice. The emitted stream could diverge from force-sync execution even though the request appeared to recover. Retiring the wrong number of speculative appends could also leave the pool offset and KV length inconsistent with the committed prefix. + +## 2. Technical Review + +The change confines evaluation errors to existing request-level error and teardown paths. Scheduler read failures record an evaluation failure and enter the established finishing teardown before retrying through guarded synchronous execution. This scheduler branch is reviewed in code, as the core test-only fault seam is not available to the server crate. + +Token waits use per-array `try_eval`, preserving overlap with the next in-flight forward. Feedback snapshots copy only mirostat/adaptive-p state, not penalty history caches. The fused stochastic path adds a host-side snapshot allocation and key-array ownership copy; it adds no evaluation or scheduling operation. No successful-path throughput result is claimed in this draft. + +Existing infallible public sampler entrypoints retain their signatures and wrap the fallible core with the documented panic message. Engine sampling uses the new fallible entrypoints. There are no new external dependencies or data migrations. CUDA is the locally validated backend; Metal and ROCm runtime validation is not claimed. + +## 3. Technical Decisions + +### Guard the actual host-read boundary + +`try_tokens_to_host` and `try_read_token` first perform the fallible per-array wait. Adaptive-p and mirostat guard every required evaluation/copy and return errors before publishing feedback updates. `DrawError::Mask` retains structured-output classification, while `DrawError::Eval` becomes an evaluation error. + +Eagerly evaluating all logits before sampling was rejected because it would serialize penalty and DRY draws behind the forward. It also would not cover every synchronous sampler entrypoint. + +### Commit pending state before fallback + +Both B=1 pipelines share their normal read-and-commit helper with the failed-submit branch. That branch unwinds only the failed submission's append, commits `pending` with no `next`, and enters synchronous fallback only if generation continues. EOS and callback stops therefore use the same stopping and cache-retirement rules as normal iterations. A readback failure retires the pending append plus any existing next append. + +### Restore only the failed fused draw's RNG advancement + +Fused graph construction consumes a random key before `schedule`. Moving the injected fault earlier would conceal the real ordering defect; adding another schedule would risk successful-path throughput. The selected implementation snapshots `random::KeySequence::default_()` after forward construction and restores it only when the existing submission schedule fails. Successful submissions keep their key advancement. + +This relies on the public copy construction/assignment of `KeySequence` in MLX pin `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a`. The CXX snapshot owns its copied key array; its lifetime stays inside one synchronous call on the owning thread. A regression verifies restoration and thread isolation. Per-row failed draws retain the documented limitation that a discarded stochastic draw can shift the seeded stream by its consumed key. + +## 4. Implementation Details + +| Area | Result | +|---|---| +| Fused, per-row, plain-drafter and scheduler token reads | Backend read errors follow fallible request-error paths | +| Mirostat v1/v2/greedy and adaptive-p | All required sampler host reads propagate errors | +| In-loop submit failure | Retire failed append, commit pending token, then conditionally fall back | +| Row draw failure | Restore mirostat/adaptive-p feedback before caller teardown | +| Fused RNG | Restore failed submission's thread-local key without an extra schedule | +| Fault seam | One-shot, thread-local Submit/Draw/Readback/SamplerRead plans; constant `Ok(())` outside tests | + +The production fault seam is `#[inline(always)]` and returns `Ok(())` under `cfg(not(test))`; it has no runtime fault plan. This conclusion is based on source configuration, not release-binary disassembly. + +## 5. Learning Points + +A lazy graph can mutate sampler state before GPU execution succeeds. Correct recovery must distinguish graph construction, scheduling, readback, and token commitment. A final output comparison alone can miss a discarded greedy draw, because redrawing may produce the same token; seeded failure tests and targeted mutations expose that gap. + +## 7. Change Summary + +Before report additions, the branch changes 16 files with 868 insertions and 116 deletions relative to `24e7e7ff`. Implementation commits are `6f3c384b` (readback/state handling), `4669eeec` (seeded stopping boundaries), and `7facb4fa` (sampler output type aliases). + +## 8. Follow-up Actions + +- Confirm PR CI results before central merge. +- Keep the separate plain-drafter failed-submit/retraction problem outside this change, as specified by #2258. + +## Appendix: Validation Recorded So Far + +| Check | Observed result | +|---|---| +| Formatting and diff checks | Passed | +| Core CUDA lib/tests clippy, `-D warnings` | Passed after factoring three complex return types into aliases | +| Core CUDA test-fast test build | Passed | +| Original engine fault suite | 11 passed, GPU lock held, one test thread | +| Plain-drafter readback fault | 1 passed | +| Thread-local fault-plan regression | 1 passed | +| Remove pending commitment | Four separate seeded fused/mirostat/adaptive/penalty stream tests failed as expected | +| Remove feedback restoration | Feedback regression failed as expected | +| Remove fused RNG restoration | Seeded fused regression failed as expected | +| Restored implementation after mutations | All 13 fault regressions passed again, including seeded EOS/callback tests | +| New seeded EOS/callback tests | Both passed in the engine suite; each failed when pending commitment was removed independently from either pipeline (four negative checks) | +| Full engine/sampling/server suites and contracts | Engine 63 passed; sampling 202 passed; server batch 526 passed, 9 ignored; three contracts passed | +| Successful-path throughput measurement | Not measured for this draft | + +Focused evidence is in `/tmp/test2258-final.log`, `/tmp/plain2258.log`, `/tmp/seam2258.log`, and `/tmp/mutation2258-{discard-pending,feedback,fused-rng}.log`. Central verification writes `/tmp/mlxcel-auto-plan/2258-*.log`. Stopping-boundary mutations are recorded in `/tmp/mutation2258-stops-summary.log` and `/tmp/mutation2258-pending-{eos,callback}-{fused,row}.log`; the final positive rerun is `/tmp/test2258-restored-stops.log`. These are session-local evidence paths, not committed artifacts. + +Related work: issue #2258; request-local backend error contract #822; per-row pipeline #2256 and #2229. diff --git a/TECHNICAL_REPORTS/2269-decode-readback-sampler-recovery-20261011.ko.md b/TECHNICAL_REPORTS/2269-decode-readback-sampler-recovery-20261011.ko.md new file mode 100644 index 000000000..d123ee7c1 --- /dev/null +++ b/TECHNICAL_REPORTS/2269-decode-readback-sampler-recovery-20261011.ko.md @@ -0,0 +1,95 @@ +# 기술 보고서: PR #2269 — 디코드 읽기 오류 격리와 샘플러 상태 보존 + +**작성일**: 2026-10-11 + +**상태**: 구현 및 필수 로컬 검증 완료. 보고서 작성 시점에 CI는 진행 중이다. + +**언어**: Rust, C++ + +**위험도**: 높음. 지연 GPU 그래프, CXX 오류 경계, 추측 실행 KV 정리, 변경 가능한 샘플러/RNG 상태를 함께 다룬다. + +## 요약 + +디코드의 호스트 읽기가 실패를 반환할 수 있는 CXX 진입점을 통해 기다리도록 바꿨다. 백엔드 평가 오류가 실패를 표현하지 못하는 브리지를 넘어가는 대신 해당 요청의 오류로 전달된다. 미리 실행한 제출이 실패하면 이미 추출한 pending 토큰을 커밋한 뒤 동기 실행으로 전환하며, 행별 추출 실패 시 mirostat와 adaptive-p 피드백을 복원한다. + +시드가 있는 fused 실행의 동등성을 유지하려면 실패한 다음 제출을 구성하면서 소비한 RNG 키도 복원해야 한다. 고정 MLX 버전의 스레드별 키 시퀀스를 스냅샷으로 보관해 GPU 평가나 스케줄 호출을 추가하지 않고 복구한다. 집중 오류 테스트, 전체 검증 및 새로 강화한 종료 경계 변이 검증은 통과했으며, 복원한 구현은 오류 회귀 테스트 13개를 모두 재통과했다. + +## 1. 문제 정의 + +비동기 백엔드 오류는 호스트가 토큰, 선택 확률, 정렬된 샘플러 분포를 기다릴 때 처음 드러날 수 있다. 일부 디코드 경로는 실패 가능한 대기 전에 `Result`가 아닌 CXX 읽기를 호출해 C++ 예외가 프로세스를 종료할 수 있었다. 최종 동기 토큰 읽기만 보호하면 샘플링 내부에서 먼저 읽는 mirostat와 adaptive-p는 보호되지 않는다. + +추측 실행 실패 후 상태가 중복 진행되는 문제도 있었다. pending 토큰을 버리고 다시 추출하면 무작위 키를 한 번 더 소비하거나 피드백을 두 번 진행할 수 있다. 요청이 복구된 것처럼 보여도 force-sync 실행과 출력 스트림이 달라질 수 있었다. 추측 실행으로 추가한 KV 항목을 잘못된 개수만큼 정리하면 풀 오프셋과 KV 길이도 커밋된 접두 구간과 어긋날 수 있다. + +## 2. 기술적 검토 사항 + +평가 오류는 기존 요청 단위 오류와 정리 경로로 전달한다. 스케줄러 읽기 실패는 평가 실패를 기록하고 기존 종료 정리 경로에 진입한 다음 보호된 동기 실행으로 재시도한다. 코어의 테스트 전용 오류 주입 기능은 서버 크레이트에서 사용할 수 없으므로 이 스케줄러 분기는 코드 검토로 확인한다. + +토큰 대기는 배열별 `try_eval`을 사용해 다음 forward와의 중첩을 유지한다. 피드백 스냅샷은 mirostat/adaptive-p 상태만 복사하며 페널티 이력 캐시는 복사하지 않는다. fused 확률적 경로에는 호스트 스냅샷 할당과 키 배열 소유권 복사가 추가되지만 평가나 스케줄 호출은 추가되지 않는다. 이 초안은 정상 경로 처리량 측정 결과를 주장하지 않는다. + +기존 infallible 공개 샘플러 진입점은 시그니처를 유지하고 정해진 panic 메시지로 fallible 코어를 감싼다. 엔진 샘플링은 새 fallible 진입점을 사용한다. 외부 의존성이나 데이터 마이그레이션은 추가하지 않았다. 로컬에서 검증한 백엔드는 CUDA이며 Metal 및 ROCm 실행 검증을 수행했다고 주장하지 않는다. + +## 3. 기술적 선택과 그 이유 + +### 실제 호스트 읽기 경계 보호 + +`try_tokens_to_host`와 `try_read_token`은 먼저 실패 가능한 배열별 대기를 수행한다. Adaptive-p와 mirostat는 필요한 모든 평가 및 복사를 보호하고 피드백 변경을 게시하기 전에 오류를 반환한다. `DrawError::Mask`는 구조화 출력 오류 분류를 유지하고 `DrawError::Eval`은 평가 오류로 전달한다. + +샘플링 전에 모든 logits를 평가하는 방법은 페널티와 DRY 추출을 forward 뒤로 직렬화하므로 선택하지 않았다. 이 방법으로 모든 동기 샘플러 진입점을 보호할 수도 없다. + +### 동기 전환 전에 pending 상태 커밋 + +두 B=1 파이프라인은 정상 반복과 제출 실패 분기에서 같은 읽기 및 커밋 헬퍼를 사용한다. 실패 분기는 해당 제출이 추가한 항목만 되돌리고 `next` 없이 pending을 커밋한 뒤, 생성이 계속될 때만 동기 실행으로 전환한다. 따라서 EOS와 콜백 중지는 정상 반복과 같은 종료 및 캐시 정리 규칙을 따른다. 읽기 실패 시 pending 항목과 이미 존재하는 next 항목을 함께 정리한다. + +### 실패한 fused 추출의 RNG 진행만 복원 + +Fused 그래프 구성은 `schedule` 전에 무작위 키를 소비한다. 주입 오류 위치를 앞으로 옮기면 실제 순서 문제를 숨기고, 스케줄 호출을 하나 더 넣으면 정상 경로 처리량에 영향을 줄 수 있다. 선택한 구현은 forward 구성 후 `random::KeySequence::default_()`를 스냅샷으로 보관하고 기존 제출 스케줄이 실패한 경우에만 복원한다. 성공한 제출은 키 진행을 유지한다. + +이 구현은 MLX pin `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a`의 `KeySequence`가 제공하는 공개 복사 생성과 대입에 의존한다. CXX 스냅샷은 복사한 키 배열을 소유하며, 수명은 해당 스레드의 동기 호출 하나 안에 한정된다. 회귀 테스트는 복원과 스레드 격리를 확인한다. 행별 추출 실패에는 버린 확률적 추출이 소비한 키만큼 시드 스트림이 이동할 수 있다는 기존 문서화된 제한이 남는다. + +## 4. 구현 상세 + +| 영역 | 결과 | +|---|---| +| Fused, 행별, plain-drafter, 스케줄러 토큰 읽기 | 백엔드 읽기 오류를 fallible 요청 오류 경로로 전달 | +| Mirostat v1/v2/greedy 및 adaptive-p | 필요한 모든 샘플러 호스트 읽기에서 오류 전달 | +| 반복 중 제출 실패 | 실패한 추가 항목 정리, pending 토큰 커밋, 조건부 동기 전환 | +| 행별 추출 실패 | 호출 측 정리 전에 mirostat/adaptive-p 피드백 복원 | +| Fused RNG | 추가 스케줄 없이 실패한 제출의 스레드별 키 복원 | +| 오류 주입 | 스레드별 일회성 Submit/Draw/Readback/SamplerRead 계획; 테스트 외에는 상수 `Ok(())` | + +프로덕션 오류 주입 진입점은 `#[inline(always)]`이며 `cfg(not(test))`에서 `Ok(())`를 반환하므로 실행 시 오류 계획이 없다. 이는 소스 설정으로 확인한 결과이며 릴리스 바이너리 역어셈블리 결과는 아니다. + +## 5. 학습 포인트 + +지연 그래프는 GPU 실행 성공 전에 샘플러 상태를 변경할 수 있다. 올바른 복구를 위해 그래프 구성, 스케줄, 읽기, 토큰 커밋을 구분해야 한다. 최종 출력 비교만으로는 버린 greedy 추출을 발견하지 못할 수 있다. 다시 추출해도 같은 토큰이 나올 수 있기 때문이다. 시드 기반 실패 테스트와 대상 변이 검증으로 이 공백을 확인할 수 있다. + +## 7. 변경 요약 + +보고서 추가 전 기준으로 `24e7e7ff` 대비 16개 파일, 868줄 추가, 116줄 삭제다. 구현 커밋은 `6f3c384b`(읽기 및 상태 처리), `4669eeec`(시드 기반 종료 경계), `7facb4fa`(샘플러 출력 타입 별칭)다. + +## 8. 후속 조치 + +- 중앙 병합 전에 PR CI 결과를 확인한다. +- #2258에서 지정한 대로 별도의 plain-drafter 제출 실패 및 draft 철회 문제는 이번 변경 범위 밖에 둔다. + +## 부록: 현재까지 기록한 검증 + +| 검사 | 확인한 결과 | +|---|---| +| 포맷 및 diff 검사 | 통과 | +| 코어 CUDA lib/tests clippy, `-D warnings` | 복잡한 반환 타입 세 곳을 별칭으로 정리한 뒤 통과 | +| 코어 CUDA test-fast 테스트 빌드 | 통과 | +| 기존 엔진 오류 테스트 묶음 | GPU 잠금과 단일 테스트 스레드로 11개 통과 | +| Plain-drafter 읽기 오류 | 1개 통과 | +| 스레드별 오류 계획 회귀 | 1개 통과 | +| Pending 커밋 제거 | 별도 시드 fused/mirostat/adaptive/penalty 스트림 테스트 4개가 예상대로 실패 | +| 피드백 복원 제거 | 피드백 회귀 테스트가 예상대로 실패 | +| Fused RNG 복원 제거 | 시드 fused 회귀 테스트가 예상대로 실패 | +| 변이 후 구현 복원 | 시드 EOS/콜백 테스트를 포함한 오류 회귀 테스트 13개 모두 재통과 | +| 새 시드 EOS/콜백 테스트 | 엔진 테스트에서 둘 다 통과; 각 파이프라인의 pending 커밋을 독립적으로 제거했을 때 각각 실패(음성 검증 4건) | +| 전체 엔진/샘플링/서버 테스트와 계약 검사 | 엔진 63개 통과, 샘플링 202개 통과, 서버 배치 526개 통과 및 9개 무시, 계약 검사 3개 통과 | +| 정상 경로 처리량 측정 | 이 초안에서 측정하지 않음 | + +집중 검증 자료는 `/tmp/test2258-final.log`, `/tmp/plain2258.log`, `/tmp/seam2258.log`, `/tmp/mutation2258-{discard-pending,feedback,fused-rng}.log`에 있다. 중앙 검증은 `/tmp/mlxcel-auto-plan/2258-*.log`에 기록한다. 종료 경계 변이는 `/tmp/mutation2258-stops-summary.log`와 `/tmp/mutation2258-pending-{eos,callback}-{fused,row}.log`, 마지막 정상 재실행은 `/tmp/test2258-restored-stops.log`에 기록했다. 이 경로들은 세션 내 검증 자료이며 커밋된 산출물이 아니다. + +관련 작업: 이슈 #2258, 요청별 백엔드 오류 계약 #822, 행별 파이프라인 #2256 및 #2229. diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 70cf07778..6ad05a33d 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -51,6 +51,15 @@ namespace mlx_cxx { using namespace mlx::core; +std::unique_ptr rng_snapshot() { + return std::make_unique(); +} + +void restore_rng_snapshot(const MlxRngSnapshot& snapshot) { + random::KeySequence::default_() = snapshot.inner; +} + + namespace { using CompiledArrayFn = diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index c5f364511..d4850ffde 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -12,6 +12,16 @@ namespace mlx_cxx { +// A copy of the pinned MLX KeySequence retains only its lazy key array. +// Its public copy/assignment operations let failed fused submits restore the +// calling thread's RNG without evaluating a graph or adding a GPU schedule. +struct MlxRngSnapshot { + mlx::core::random::KeySequence inner; + MlxRngSnapshot() : inner(mlx::core::random::KeySequence::default_()) {} +}; +std::unique_ptr rng_snapshot(); +void restore_rng_snapshot(const MlxRngSnapshot& snapshot); + // Opaque wrapper struct to hold mlx::core::array // This allows cxx to manage the lifetime without exposing the complex internals struct MlxArray { diff --git a/src/lib/mlxcel-core/src/engine/direct_decode.rs b/src/lib/mlxcel-core/src/engine/direct_decode.rs index 58df6899e..d2ab78bc6 100644 --- a/src/lib/mlxcel-core/src/engine/direct_decode.rs +++ b/src/lib/mlxcel-core/src/engine/direct_decode.rs @@ -56,7 +56,7 @@ use super::direct::{BareHooks, DirectEngine, DirectEngineError, delivers_to_callback}; use super::lookahead::lookahead_feedback_input; -use super::rows::{FusedRowGate, shared_fused_params, tokens_to_host}; +use super::rows::{FusedRowGate, shared_fused_params, try_tokens_to_host}; use super::{EngineError, RowOutcome, StepBatch, StepRow}; use crate::cache::{SequenceId, SequenceStateBackend}; use crate::generate::{LanguageModel, SamplingConfig}; @@ -311,31 +311,15 @@ impl DirectEngine { match self.submit_lookahead(id, state.sampling, &input, params) { Ok(tokens) => Some(tokens), Err(err) => { - // The pending forward's append is still uncommitted. - drop(pending); - self.unwind_failed_submit(id, &err, 1, teardown)?; + self.unwind_failed_submit(id, &err, 0, teardown)?; + if self.commit_fused_pending(state, pending, None, teardown, on_token)? { + return Ok(()); + } return self.decode_sync(state, on_token); } } }; - // The sync point: waits for `pending` only, while `next` runs. - let host = tokens_to_host(&pending); - drop(pending); - let before = state.generated.len(); - let outcome = single(self.engine_mut().finish_rows(&host, &mut [state.row()]))?; - if let Some(error) = outcome.error { - // Nothing was committed: undo the read step's append and the - // next one, so the closed sequence holds no speculative state. - let n = 1 + i32::from(next.is_some()); - self.retire(id, next.as_deref(), n, teardown)?; - return Err(DirectEngineError::Row(error.message().to_string())); - } - let stopped = delivers_to_callback(&state.generated, before, &outcome) - && !on_token(outcome.token); - if stopped || outcome.finish.is_some() { - if next.is_some() { - self.retire(id, next.as_deref(), 1, teardown)?; - } + if self.commit_fused_pending(state, pending, next.as_deref(), teardown, on_token)? { return Ok(()); } match next { @@ -348,6 +332,44 @@ impl DirectEngine { } } + /// Read and commit the pending draw, shared by normal and failed submits. + /// Returns whether the run finished; `next` owns only an uncommitted append. + fn commit_fused_pending bool>( + &mut self, + state: &mut DecodeState<'_>, + pending: UniquePtr, + next: Option<&MlxArray>, + teardown: Teardown, + on_token: &mut F, + ) -> Result { + let id = state.id; + // Wait only for pending, preserving overlap with the next forward. + let host = match try_tokens_to_host(&pending) { + Ok(host) => host, + Err(msg) => { + drop(pending); + self.retire(id, next, 1 + i32::from(next.is_some()), teardown)?; + return Err(DirectEngineError::Row(msg)); + } + }; + drop(pending); + let before = state.generated.len(); + let outcome = single(self.engine_mut().finish_rows(&host, &mut [state.row()]))?; + if let Some(error) = outcome.error { + self.retire(id, next, 1 + i32::from(next.is_some()), teardown)?; + return Err(DirectEngineError::Row(error.message().to_string())); + } + let stopped = + delivers_to_callback(&state.generated, before, &outcome) && !on_token(outcome.token); + if stopped || outcome.finish.is_some() { + if next.is_some() { + self.retire(id, next, 1, teardown)?; + } + return Ok(true); + } + Ok(false) + } + /// Submit one pipelined step for `id` on `input` (`[1, 1]`). pub(super) fn submit_lookahead( &mut self, diff --git a/src/lib/mlxcel-core/src/engine/direct_decode_rows.rs b/src/lib/mlxcel-core/src/engine/direct_decode_rows.rs index 639efe77d..c2f644be4 100644 --- a/src/lib/mlxcel-core/src/engine/direct_decode_rows.rs +++ b/src/lib/mlxcel-core/src/engine/direct_decode_rows.rs @@ -29,6 +29,12 @@ //! synchronous step's, and forwards consume no randomness, so greedy and //! seeded streams equal [`super::DirectEngine::decode_sync`]'s token for token. //! +//! A failed draw restores sampler feedback before synchronous fallback. A +//! stochastic draw already consumed a thread-local RNG key while building its +//! graph. This per-row path does not retain an RNG snapshot, so a discarded +//! draw's key is not restored and the failure may shift a seeded stream by one +//! draw. A failed forward submit commits the pending draw first. +//! //! The overlap differs by sampler. Penalties, DRY and the extended chain //! without adaptive-p keep the draw lazy on the device, so the draw and the //! finish step both overlap the next forward. Mirostat and adaptive-p read the @@ -47,7 +53,7 @@ use super::lookahead::lookahead_feedback_input; use super::{EngineError, StepBatch}; use crate::cache::SequenceId; use crate::generate::LanguageModel; -use crate::{MlxArray, UniquePtr, ffi}; +use crate::{MlxArray, UniquePtr}; impl DirectEngine { /// The per-row lookahead pipeline (see the module docs). Holds at most @@ -86,45 +92,15 @@ impl DirectEngine { match self.submit_row_forward(id, &input) { Ok(logits) => Some(logits), Err(err) => { - // The pending draw's forward is still uncommitted. - drop(pending); - self.unwind_failed_submit(id, &err, 1, teardown)?; + self.unwind_failed_submit(id, &err, 0, teardown)?; + if self.commit_row_pending(state, pending, None, teardown, on_token)? { + return Ok(()); + } return self.decode_sync(state, on_token); } } }; - // The sync point: waits for `pending` only, while `next` runs. - // `try_eval` is the per-array wait (it enqueues nothing behind - // `next`) behind the fallible boundary, as in the synchronous - // step: a backend throw surfacing at the wait fails the run - // instead of aborting the process in the infallible readback - // (#822). `item_i32` then reads the per-row draw in whatever - // integer dtype the chain returned. - if let Err(err) = crate::try_eval(&pending) { - drop(pending); - // Nothing was committed: undo the read step's append and the - // next one, as after a row error. - let n = 1 + i32::from(next.is_some()); - self.retire(id, next.as_deref(), n, teardown)?; - return Err(DirectEngineError::Row(err.to_string())); - } - let token = ffi::item_i32(&pending); - drop(pending); - let before = state.generated.len(); - let outcome = single(self.engine_mut().finish_rows(&[token], &mut [state.row()]))?; - if let Some(error) = outcome.error { - // Nothing was committed: undo the read step's append and the - // next one, as the fused pipeline does. - let n = 1 + i32::from(next.is_some()); - self.retire(id, next.as_deref(), n, teardown)?; - return Err(DirectEngineError::Row(error.message().to_string())); - } - let stopped = delivers_to_callback(&state.generated, before, &outcome) - && !on_token(outcome.token); - if stopped || outcome.finish.is_some() { - if next.is_some() { - self.retire(id, next.as_deref(), 1, teardown)?; - } + if self.commit_row_pending(state, pending, next.as_deref(), teardown, on_token)? { return Ok(()); } // The budget check above and the finish step's length rule agree, @@ -144,6 +120,44 @@ impl DirectEngine { } } + /// Read and commit the pending draw, shared by normal and failed submits. + /// Returns whether the run finished; `next` owns only an uncommitted append. + fn commit_row_pending bool>( + &mut self, + state: &mut DecodeState<'_>, + pending: UniquePtr, + next: Option<&MlxArray>, + teardown: Teardown, + on_token: &mut F, + ) -> Result { + let id = state.id; + // Wait only for pending, preserving overlap with the next forward. + let host = match super::rows::try_read_token(&pending) { + Ok(host) => host, + Err(msg) => { + drop(pending); + self.retire(id, next, 1 + i32::from(next.is_some()), teardown)?; + return Err(DirectEngineError::Row(msg)); + } + }; + drop(pending); + let before = state.generated.len(); + let outcome = single(self.engine_mut().finish_rows(&[host], &mut [state.row()]))?; + if let Some(error) = outcome.error { + self.retire(id, next, 1 + i32::from(next.is_some()), teardown)?; + return Err(DirectEngineError::Row(error.message().to_string())); + } + let stopped = + delivers_to_callback(&state.generated, before, &outcome) && !on_token(outcome.token); + if stopped || outcome.finish.is_some() { + if next.is_some() { + self.retire(id, next, 1, teardown)?; + } + return Ok(true); + } + Ok(false) + } + /// Submit the forward of one per-row pipelined step for `id` on `input` /// (`[1, 1]`). fn submit_row_forward( diff --git a/src/lib/mlxcel-core/src/engine/direct_decode_rows_tests.rs b/src/lib/mlxcel-core/src/engine/direct_decode_rows_tests.rs index 4f566ce98..e50095f7f 100644 --- a/src/lib/mlxcel-core/src/engine/direct_decode_rows_tests.rs +++ b/src/lib/mlxcel-core/src/engine/direct_decode_rows_tests.rs @@ -35,7 +35,7 @@ const NOISY_VOCAB: i32 = 16; /// token, spread over `[0, 3)` with small gaps, so a history penalty, DRY or /// a stale history changes the draw. #[derive(Default)] -struct NoisyModel { +pub(super) struct NoisyModel { forwards: RefCell, } diff --git a/src/lib/mlxcel-core/src/engine/eval_fault_tests.rs b/src/lib/mlxcel-core/src/engine/eval_fault_tests.rs new file mode 100644 index 000000000..649657762 --- /dev/null +++ b/src/lib/mlxcel-core/src/engine/eval_fault_tests.rs @@ -0,0 +1,395 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Decode fault regressions: backend errors stay request-local and speculative +//! state agrees with the committed synchronous prefix. +use super::direct_decode_rows_tests::NoisyModel; +use super::direct_decode_tests::{CountModel, OwnedModel, trace}; +use super::*; +use crate::eval_faults::{FaultGuard, FaultPoint}; +use crate::from_slice_i32; +use crate::generate::SamplingConfig; +use crate::sampling::{LogprobSource, LogprobsConfig}; +use crate::sampling_row_step::RowSampler; + +fn configs() -> Vec { + let base = SamplingConfig { + temperature: 0.8, + seed: Some(7), + ..SamplingConfig::default() + }; + vec![ + base.clone(), + SamplingConfig { + mirostat: 2, + ..base.clone() + }, + SamplingConfig { + adaptive_target: 0.5, + ..base.clone() + }, + SamplingConfig { + repetition_penalty: 1.3, + ..base + }, + ] +} + +fn check_seeded_submit_stream(mut sampling: SamplingConfig) { + sampling.token_bias.insert(15, f32::NEG_INFINITY); + let model = NoisyModel::default(); + let mut sync = DirectEngine::new(&model, 0).with_force_sync(true); + let expected = trace(&mut sync, &[1, 4], 32, &sampling, None); + assert_eq!(expected.tokens.len(), 32); + for nth in [0, 1, 5] { + let mut piped = DirectEngine::new(&model, 0).with_force_sync(false); + let fault = FaultGuard::arm(FaultPoint::Submit, nth); + let actual = trace(&mut piped, &[1, 4], 32, &sampling, None); + assert!(fault.fired(), "the request must reach the failed submit"); + assert_eq!(actual, expected, "config {sampling:?}, submit {nth}"); + } +} + +#[test] +fn failed_in_loop_submits_preserve_seeded_fused_stream() { + check_seeded_submit_stream(configs()[0].clone()); +} + +#[test] +fn failed_in_loop_submits_preserve_seeded_mirostat_stream() { + check_seeded_submit_stream(configs()[1].clone()); +} + +#[test] +fn failed_in_loop_submits_preserve_seeded_adaptive_stream() { + check_seeded_submit_stream(configs()[2].clone()); +} + +#[test] +fn failed_in_loop_submits_preserve_seeded_penalty_stream() { + check_seeded_submit_stream(configs()[3].clone()); +} + +#[test] +fn failed_submit_commits_pending_eos_and_callback_stop() { + for sampling in [ + SamplingConfig::greedy(), + SamplingConfig { + repetition_penalty: 1.3, + ..SamplingConfig::greedy() + }, + ] { + for (prompt, deliveries) in [(&[5][..], None), (&[0][..], Some(2))] { + let model = CountModel::default(); + let mut sync = DirectEngine::new(&model, 0).with_force_sync(true); + let expected = trace(&mut sync, prompt, 16, &sampling, deliveries); + let mut piped = DirectEngine::new(&model, 0).with_force_sync(false); + let _fault = FaultGuard::arm(FaultPoint::Submit, 1); + assert_eq!( + trace(&mut piped, prompt, 16, &sampling, deliveries), + expected + ); + } + } +} + +#[test] +fn failed_submit_commits_seeded_pending_eos() { + for mut sampling in [configs()[0].clone(), configs()[3].clone()] { + sampling.token_bias.insert(15, f32::NEG_INFINITY); + let model = NoisyModel::default(); + let mut sync = DirectEngine::new(&model, 0).with_force_sync(true); + let baseline = trace(&mut sync, &[1, 4], 16, &sampling, None); + // Choose an EOS reached in-loop without assuming backend-specific + // categorical samples. A greedy redraw could conceal a lost pending draw. + let eos_index = (2..baseline.tokens.len() - 1) + .find(|&i| !baseline.tokens[..i].contains(&baseline.tokens[i])) + .expect("the seeded stream must introduce an in-loop stop token"); + sampling.stop_token_ids = vec![baseline.tokens[eos_index]]; + let expected = trace(&mut sync, &[1, 4], 16, &sampling, None); + assert_eq!(expected.tokens, baseline.tokens[..eos_index]); + let mut piped = DirectEngine::new(&model, 0).with_force_sync(false); + // Submit zero primes the pipeline; submit eos_index fails while the + // already drawn EOS is pending, before its readback and commitment. + let fault = FaultGuard::arm(FaultPoint::Submit, eos_index); + let actual = trace(&mut piped, &[1, 4], 16, &sampling, None); + assert!( + fault.fired(), + "the pending EOS must reach the failed submit" + ); + assert_eq!(actual, expected, "config {sampling:?}"); + } +} + +#[test] +fn failed_submit_commits_seeded_pending_callback_stop() { + for mut sampling in [configs()[0].clone(), configs()[3].clone()] { + sampling.token_bias.insert(15, f32::NEG_INFINITY); + let model = NoisyModel::default(); + let mut sync = DirectEngine::new(&model, 0).with_force_sync(true); + let expected = trace(&mut sync, &[1, 4], 16, &sampling, Some(2)); + assert_eq!(expected.delivered.len(), 2); + let mut piped = DirectEngine::new(&model, 0).with_force_sync(false); + let fault = FaultGuard::arm(FaultPoint::Submit, 1); + let actual = trace(&mut piped, &[1, 4], 16, &sampling, Some(2)); + assert!( + fault.fired(), + "the callback stop must reach the failed submit" + ); + assert_eq!(actual, expected, "config {sampling:?}"); + } +} + +#[test] +fn readback_helpers_return_errors() { + let tokens = from_slice_i32(&[2], &[1]); + { + let _fault = FaultGuard::arm(FaultPoint::Readback, 0); + assert_eq!( + try_tokens_to_host(&tokens), + Err("injected eval failure".into()) + ); + } + let _fault = FaultGuard::arm(FaultPoint::Readback, 0); + assert_eq!( + super::rows::try_read_token(&tokens), + Err("injected eval failure".into()) + ); +} + +#[test] +fn failed_readbacks_unwind_to_the_committed_prefix() { + for sampling in [ + SamplingConfig::greedy(), + SamplingConfig { + repetition_penalty: 1.3, + ..SamplingConfig::greedy() + }, + ] { + for nth in [0, 2] { + let model = CountModel::default(); + let mut piped = DirectEngine::new(&model, 0).with_force_sync(false); + let id = piped.open_sequence().unwrap(); + let mut delivered = Vec::new(); + { + let _fault = FaultGuard::arm(FaultPoint::Readback, nth); + let result = piped.run_open_sequence( + id, + &DirectRequest::text(&[0], 16, &sampling), + &sampling, + 0, + |t| { + delivered.push(t); + true + }, + ); + assert!( + matches!(result, Err(DirectEngineError::Row(ref msg)) if msg == "injected eval failure") + ); + } + let set = piped.engine().pool().get(id).unwrap(); + let actual = (set.current_offset, set.caches[0].seq_len()); + let mut sync = DirectEngine::new(&model, 0).with_force_sync(true); + let expected = trace(&mut sync, &[0], delivered.len(), &sampling, None); + assert_eq!(delivered, expected.tokens); + assert_eq!(actual.1, expected.kv_len.unwrap()); + assert_eq!(actual.0, expected.offset); + piped.close_sequence(id); + assert!(piped.engine().pool().get(id).is_none()); + } + } +} + +#[test] +fn discard_failures_release_the_sequence() { + for sampling in [ + SamplingConfig::greedy(), + SamplingConfig { + repetition_penalty: 1.3, + ..SamplingConfig::greedy() + }, + ] { + let mut faults = vec![(FaultPoint::Submit, 1)]; + if sampling.repetition_penalty != 1.0 { + faults.extend([(FaultPoint::Draw, 0), (FaultPoint::Draw, 1)]); + } + for (point, nth) in faults { + let model = OwnedModel::new(false); + let mut client = DirectEngine::new(&model, 0).with_force_sync(false); + let _fault = FaultGuard::arm(point, nth); + assert!(matches!( + client.run(&[0], 16, &sampling), + Err(DirectEngineError::Step(EngineError::Eval(_))) + )); + assert_eq!(model.released.borrow().len(), 1); + } + } +} + +#[test] +fn sampler_reads_report_eval_errors_and_failed_draws_restore_feedback() { + let configs = [ + SamplingConfig { + mirostat: 1, + temperature: 0.8, + ..SamplingConfig::default() + }, + SamplingConfig { + mirostat: 2, + temperature: 0.8, + ..SamplingConfig::default() + }, + SamplingConfig { + mirostat: 2, + ..SamplingConfig::greedy() + }, + SamplingConfig { + adaptive_target: 0.5, + temperature: 0.8, + ..SamplingConfig::default() + }, + ]; + for sampling in configs { + // Fail every evaluation/copy boundary, including later reads. + let reads = if sampling.mirostat == 1 { + 3 + } else if sampling.adaptive_target >= 0.0 { + 2 + } else { + 1 + }; + for point in [FaultPoint::SamplerRead, FaultPoint::Draw] { + for nth in 0..if point == FaultPoint::Draw { 1 } else { reads } { + let mut engine = Engine::with_capacity(CountModel::default(), 1); + let id = engine.open(SequenceSpec::default()).unwrap(); + let input = from_slice_i32(&[3], &[1, 1]); + let logits = engine + .submit_forward(&StepBatch { + seq_ids: &[id], + input: &input, + }) + .unwrap(); + let mut sampler = RowSampler::new(&sampling); + // Establish noninitial feedback to ensure restoration retains it. + let initial = sampler.draw_and_accept(&logits, &sampling, &[3]).unwrap(); + crate::try_eval(&initial.token).unwrap(); + let before = sampler.feedback_snapshot(); + let (mut history, mut generated) = (vec![3], vec![3]); + let logprobs = LogprobsConfig::default(); + let mut row = StepRow { + seq_id: id, + sampler: &mut sampler, + sampling: &sampling, + token_history: &mut history, + generated: &mut generated, + eos: &[7], + max_tokens: 8, + logprobs: &logprobs, + needs_mask: false, + needs_override: false, + hooks: BareHooks, + }; + { + let _fault = FaultGuard::arm(point, nth); + assert!( + matches!(engine.draw_row(&logits, &mut row), Err(EngineError::Eval(ref msg)) if msg == "injected eval failure") + ); + } + assert_eq!(row.sampler.feedback_snapshot(), before); + if point == FaultPoint::SamplerRead { + let distribution = LogprobsConfig { + enabled: true, + source: LogprobSource::PostSampling, + ..LogprobsConfig::default() + }; + for logprobs in [&logprobs, &distribution] { + row.logprobs = logprobs; + let _fault = FaultGuard::arm(point, nth); + let outcome = engine + .step( + &StepBatch { + seq_ids: &[id], + input: &input, + }, + std::slice::from_mut(&mut row), + ) + .unwrap(); + assert!( + matches!(outcome.rows[0].error, Some(RowError::Eval(ref msg)) if msg == "injected eval failure") + ); + assert_eq!(row.sampler.feedback_snapshot(), before); + } + } + engine.close(id); + } + } + } +} + +#[test] +fn rng_snapshots_restore_only_the_owning_threads_key_sequence() { + fn draw() -> Vec { + let logits = crate::ffi::zeros(&[32, 16], crate::dtype::FLOAT32); + let token = crate::ffi::fused_sample(&logits, 0.8, 0, 1.0, 0.0); + try_tokens_to_host(&token).unwrap() + } + crate::ffi::random_seed(7); + let snapshot = crate::ffi::rng_snapshot().unwrap(); + let first = draw(); + std::thread::spawn(|| { + crate::ffi::random_seed(99); + let snapshot = crate::ffi::rng_snapshot().unwrap(); + let first = draw(); + crate::ffi::restore_rng_snapshot(&snapshot); + assert_eq!(draw(), first); + }) + .join() + .unwrap(); + let second = draw(); + assert_ne!(first, second, "successful draws advance the key"); + crate::ffi::restore_rng_snapshot(&snapshot); + assert_eq!(draw(), first); + assert_eq!(draw(), second, "another thread must not change this stream"); +} + +#[test] +fn failed_draws_unwind_before_synchronous_fallback() { + fn check(model: &M, sampling: &SamplingConfig) { + let mut sync = DirectEngine::new(model, 0).with_force_sync(true); + let expected = trace(&mut sync, &[0], 16, sampling, None); + for nth in [0, 1] { + let mut piped = DirectEngine::new(model, 0).with_force_sync(false); + let fault = FaultGuard::arm(FaultPoint::Draw, nth); + let actual = trace(&mut piped, &[0], 16, sampling, None); + assert!(fault.fired()); + assert_eq!(actual, expected); + } + } + for sampling in [ + SamplingConfig { + repetition_penalty: 1.3, + ..SamplingConfig::greedy() + }, + SamplingConfig { + mirostat: 2, + ..SamplingConfig::greedy() + }, + ] { + check(&CountModel::default(), &sampling); + let model = OwnedModel::new(true); + check(&model, &sampling); + let released = model.released.borrow(); + assert!(released.iter().all(|len| *len == released[0])); + } +} diff --git a/src/lib/mlxcel-core/src/engine/mod.rs b/src/lib/mlxcel-core/src/engine/mod.rs index 6c8f72671..4a0aa587a 100644 --- a/src/lib/mlxcel-core/src/engine/mod.rs +++ b/src/lib/mlxcel-core/src/engine/mod.rs @@ -60,7 +60,7 @@ use rows::fail_batch; pub use rows::{ FusedRowGate, RowError, RowOutcome, StepRow, StepRowHooks, finish_row, fused_params, row_biases, row_logits, sample_and_finish, sample_and_finish_row, shared_fused_params, - tokens_to_host, + tokens_to_host, try_tokens_to_host, }; pub use speculative::{SpeculativeRounds, SpeculativeRun, SpeculativeRunError}; pub use verify_round::{VerifyRound, token_only_rounds_unsupported}; @@ -286,11 +286,24 @@ impl Engine { biases: &[&TokenBiasMap], ) -> Result, EngineError> { let logits = self.speculative_forward(batch)?; - // The synchronous fused path's draw (token bias in the per-row - // sampler's chain position, then the pre-fused row filters), so the - // pipelined draw samples from the identical distribution. + // Graph construction spends a thread-local RNG key before scheduling. + // Retain it for a failed fused submit without adding another async_eval. + // The snapshot never leaves this synchronous call or its owning thread. + let rng = if params.temperature > 0.0 { + Some(crate::ffi::rng_snapshot().map_err(|e| EngineError::Eval(e.to_string()))?) + } else { + None + }; let tokens = batched_fused_sample_tokens(&logits, params, biases); - schedule(&tokens)?; + let submitted = crate::eval_faults::check(crate::eval_faults::FaultPoint::Submit) + .map_err(EngineError::Eval) + .and_then(|()| schedule(&tokens)); + if let Err(err) = submitted { + if let Some(rng) = rng { + crate::ffi::restore_rng_snapshot(&rng); + } + return Err(err); + } Ok(tokens) } @@ -308,6 +321,8 @@ impl Engine { batch: &StepBatch<'_>, ) -> Result, EngineError> { let logits = self.speculative_forward(batch)?; + crate::eval_faults::check(crate::eval_faults::FaultPoint::Submit) + .map_err(EngineError::Eval)?; schedule(&logits)?; Ok(logits) } @@ -336,11 +351,21 @@ impl Engine { row.seq_id ))); } - let draw = row - .sampler - .draw_and_accept(logits, row.sampling, row.token_history); - schedule(&draw.token)?; - Ok(draw.token) + let snapshot = row.sampler.feedback_snapshot(); + let result = (|| { + let draw = row + .sampler + .draw_and_accept(logits, row.sampling, row.token_history) + .map_err(EngineError::Eval)?; + crate::eval_faults::check(crate::eval_faults::FaultPoint::Draw) + .map_err(EngineError::Eval)?; + schedule(&draw.token)?; + Ok(draw.token) + })(); + if result.is_err() { + row.sampler.restore_feedback(snapshot); + } + result } /// The collect half of a pipelined step: `tokens[i]` is the host token @@ -499,3 +524,6 @@ mod speculative_plain_tests; mod tests; #[cfg(test)] mod verify_round_tests; + +#[cfg(test)] +mod eval_fault_tests; diff --git a/src/lib/mlxcel-core/src/engine/rows.rs b/src/lib/mlxcel-core/src/engine/rows.rs index 6a1776cec..71c167292 100644 --- a/src/lib/mlxcel-core/src/engine/rows.rs +++ b/src/lib/mlxcel-core/src/engine/rows.rs @@ -33,7 +33,7 @@ use crate::sampling::{ batched_fused_sample_tokens, compute_logprobs, compute_post_sampling_probs, row_supports_fused_batch_except_bias, }; -use crate::sampling_row_step::{LogitMask, RowSampler, TokenDraw}; +use crate::sampling_row_step::{DrawError, LogitMask, RowSampler, TokenDraw}; use crate::{MlxArray, UniquePtr}; /// The caller-side hooks one row's sampling and finishing need. @@ -289,7 +289,10 @@ pub fn sample_and_finish_row( distribution, } = match draw { Ok(draw) => draw, - Err(msg) => return RowOutcome::failed(row.seq_id, RowError::Structured(msg)), + Err(DrawError::Mask(msg)) => { + return RowOutcome::failed(row.seq_id, RowError::Structured(msg)); + } + Err(DrawError::Eval(msg)) => return RowOutcome::failed(row.seq_id, RowError::Eval(msg)), }; // #822: force-evaluate the sampled token through the fallible boundary. // On an MLX throw only this row fails; the infallible readback below @@ -392,3 +395,17 @@ pub fn tokens_to_host(tokens: &MlxArray) -> Vec { pub fn row_logits(logits: &MlxArray, i: usize) -> UniquePtr { ffi::slice(logits, &[i as i32, 0, 0], &[i as i32 + 1, 1, i32::MAX]) } + +/// Wait through the fallible boundary before the overlap-preserving byte read. +pub fn try_tokens_to_host(tokens: &MlxArray) -> Result, String> { + crate::eval_faults::check(crate::eval_faults::FaultPoint::Readback)?; + crate::try_eval(tokens).map_err(|e| e.to_string())?; + Ok(tokens_to_host(tokens)) +} + +/// Read a per-row token without letting an asynchronous backend error escape. +pub(crate) fn try_read_token(token: &MlxArray) -> Result { + crate::eval_faults::check(crate::eval_faults::FaultPoint::Readback)?; + crate::try_eval(token).map_err(|e| e.to_string())?; + Ok(ffi::item_i32(token)) +} diff --git a/src/lib/mlxcel-core/src/engine/speculative_plain.rs b/src/lib/mlxcel-core/src/engine/speculative_plain.rs index c8d8d7a1b..f0ef4421e 100644 --- a/src/lib/mlxcel-core/src/engine/speculative_plain.rs +++ b/src/lib/mlxcel-core/src/engine/speculative_plain.rs @@ -24,7 +24,7 @@ use super::direct::{DirectEngine, DirectEngineError, delivers_to_callback}; use super::direct_decode::{DecodeState, Teardown, single}; use super::lookahead::lookahead_feedback_input; -use super::rows::tokens_to_host; +use super::rows::try_tokens_to_host; use super::speculative::{SpeculativeRounds, SpeculativeRunError}; use crate::drafter::Drafter; use crate::generate::LanguageModel; @@ -84,7 +84,19 @@ impl DirectEngine { if next.is_none() && asked { drafter.retract_draft(draft); } - let host = tokens_to_host(&pending); + let host = match try_tokens_to_host(&pending) { + Ok(host) => host, + Err(msg) => { + drop(pending); + self.retire( + id, + next.as_deref(), + 1 + i32::from(next.is_some()), + Teardown::Unwind, + )?; + return Err(DirectEngineError::Row(msg).into()); + } + }; let before = state.generated.len(); let outcome = single(self.engine_mut().finish_rows(&host, &mut [state.row()]))?; if let Some(error) = outcome.error { diff --git a/src/lib/mlxcel-core/src/engine/speculative_plain_tests.rs b/src/lib/mlxcel-core/src/engine/speculative_plain_tests.rs index abef2ed4f..af2c001e8 100644 --- a/src/lib/mlxcel-core/src/engine/speculative_plain_tests.rs +++ b/src/lib/mlxcel-core/src/engine/speculative_plain_tests.rs @@ -386,3 +386,46 @@ fn an_eos_inside_a_verify_block_leaves_the_offset_where_the_step_does() { let (stepped, _, _) = drafted(true, &mut plain, &[0], 20, &stop, None); assert_eq!(verified, stepped); } + +#[test] +fn failed_plain_readback_unwinds_to_the_committed_prefix() { + use crate::eval_faults::{FaultGuard, FaultPoint}; + let sampling = SamplingConfig::greedy(); + for nth in [0, 2] { + let model = CycleModel::default(); + let mut client = DirectEngine::new(&model, 0).with_force_sync(false); + let id = client.open_sequence().unwrap(); + let mut delivered = Vec::new(); + { + let _fault = FaultGuard::arm(FaultPoint::Readback, nth); + let mut drafter = ScriptedDrafter::new(&[], true); + let result = client.speculate( + id, + &DirectRequest::text(&[0], 16, &sampling), + &sampling, + &mut drafter, + 4, + |token| { + delivered.push(token); + true + }, + ); + assert!(result.is_err()); + assert!( + result + .err() + .unwrap() + .to_string() + .contains("injected eval failure") + ); + } + let set = client.engine().pool().get(id).unwrap(); + let actual = (set.current_offset, set.caches[0].seq_len()); + let mut drafter = ScriptedDrafter::new(&[], false); + let (expected, _, _) = drafted(true, &mut drafter, &[0], delivered.len(), &sampling, None); + assert_eq!(delivered, expected.tokens); + assert_eq!(actual, (expected.offset, expected.kv_len)); + client.close_sequence(id); + assert!(client.engine().pool().get(id).is_none()); + } +} diff --git a/src/lib/mlxcel-core/src/eval_faults.rs b/src/lib/mlxcel-core/src/eval_faults.rs new file mode 100644 index 000000000..5d7eea255 --- /dev/null +++ b/src/lib/mlxcel-core/src/eval_faults.rs @@ -0,0 +1,105 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Thread-local, one-shot failure injection at decode evaluation boundaries. + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FaultPoint { + Submit, + Draw, + Readback, + SamplerRead, +} + +#[cfg(not(test))] +#[inline(always)] +pub(crate) fn check(_point: FaultPoint) -> Result<(), String> { + Ok(()) +} + +#[cfg(test)] +thread_local! { + static PLAN: std::cell::Cell> = const { std::cell::Cell::new(None) }; +} + +#[cfg(test)] +pub(crate) fn check(point: FaultPoint) -> Result<(), String> { + PLAN.with(|plan| { + if let Some((armed, remaining)) = plan.get() + && armed == point + { + if remaining == 0 { + plan.set(None); + return Err("injected eval failure".into()); + } + plan.set(Some((armed, remaining - 1))); + } + Ok(()) + }) +} + +#[cfg(test)] +pub(crate) struct FaultGuard(std::marker::PhantomData>); + +#[cfg(test)] +impl FaultGuard { + pub(crate) fn fired(&self) -> bool { + PLAN.with(|plan| plan.get().is_none()) + } + + pub(crate) fn arm(point: FaultPoint, nth: usize) -> Self { + PLAN.with(|plan| { + assert!(plan.get().is_none(), "fault plan already armed"); + plan.set(Some((point, nth))); + }); + Self(std::marker::PhantomData) + } +} + +#[cfg(test)] +impl Drop for FaultGuard { + fn drop(&mut self) { + PLAN.with(|plan| plan.set(None)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn plans_are_one_shot_scoped_and_thread_local() { + { + let _guard = FaultGuard::arm(FaultPoint::Submit, 1); + assert!(check(FaultPoint::Draw).is_ok()); + std::thread::spawn(|| { + assert!(check(FaultPoint::Submit).is_ok()); + let _guard = FaultGuard::arm(FaultPoint::Readback, 0); + assert!(check(FaultPoint::Readback).is_err()); + }) + .join() + .unwrap(); + assert!(check(FaultPoint::Submit).is_ok()); + assert_eq!( + check(FaultPoint::Submit), + Err("injected eval failure".into()) + ); + assert!(check(FaultPoint::Submit).is_ok()); + } + { + let _guard = FaultGuard::arm(FaultPoint::Draw, 0); + } + assert!(check(FaultPoint::Draw).is_ok()); + } +} diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index cd7735d5d..7839e0564 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -31,6 +31,9 @@ mod ffi { /// Opaque wrapper for mlx::core::array type MlxArray; + type MlxRngSnapshot; + fn rng_snapshot() -> Result>; + fn restore_rng_snapshot(snapshot: &MlxRngSnapshot); /// Opaque wrapper for one MLX quantize() result triple. type MlxQuantizedWeights; @@ -4011,6 +4014,7 @@ pub mod prefill_plan; // every forward through `open` / `prefill` / `step` / `close`; a single // sequence is a batch of one (ADR 0007, #2172). pub mod engine; +mod eval_faults; pub mod runtime_lora; diff --git a/src/lib/mlxcel-core/src/sampling.rs b/src/lib/mlxcel-core/src/sampling.rs index af4d85534..b2612746d 100644 --- a/src/lib/mlxcel-core/src/sampling.rs +++ b/src/lib/mlxcel-core/src/sampling.rs @@ -424,10 +424,24 @@ fn sample_token_optimized_core( state: Option<&mut SamplerState>, ) -> (UniquePtr, UniquePtr) { let (token, adjusted, _probs) = - sample_token_optimized_core_full(logits, config, token_history, state, false); + sample_token_optimized_core_full(logits, config, token_history, state, false) + .unwrap_or_else(|e| panic!("sampler host read failed: {e}")); (token, adjusted) } +/// Sampled token, adjusted logits, and the post-chain probability distribution. +pub type SampleWithDistribution = ( + UniquePtr, + UniquePtr, + UniquePtr, +); + +type SampleWithOptionalDistribution = ( + UniquePtr, + UniquePtr, + Option>, +); + /// Full-form core sampler shared by every entry point: samples one token and, /// when `want_distribution` is set, also returns the float32 `[1, vocab]` /// post-chain distribution the draw came from (b10621's `post_sampling_probs` @@ -449,11 +463,7 @@ fn sample_token_optimized_core_full( token_history: &[i32], mut state: Option<&mut SamplerState>, want_distribution: bool, -) -> ( - UniquePtr, - UniquePtr, - Option>, -) { +) -> Result { if config.effective_mirostat() != 0 { return mirostat_sample_core(logits, config, state, want_distribution); } @@ -464,14 +474,14 @@ fn sample_token_optimized_core_full( let gate = xtc_gate_draw(config); let filtered = apply_extended_chain(&adjusted, config, gate.as_deref()); let (token, probs) = if config.effective_adaptive_target() >= 0.0 { - adaptive_p_sample(&filtered, config, state, want_distribution) + adaptive_p_sample(&filtered, config, state, want_distribution)? } else { let token = ffi::fused_sample(&filtered, 1.0, 0, 1.0, 0.0); let probs = want_distribution.then(|| ffi::softmax(&filtered, -1)); (token, probs) }; crate::sampling_dispatch::report_sampling_dispatch(); - return (token, adjusted, probs); + return Ok((token, adjusted, probs)); } let last_logits = preprocess_logits_for_sampling(logits, config, token_history, state); @@ -483,7 +493,7 @@ fn sample_token_optimized_core_full( crate::sampling_dispatch::report_sampling_dispatch(); let probs = want_distribution .then(|| fused_sample_probs_dispatch(&last_logits, config, gate.as_deref())); - (token, last_logits, probs) + Ok((token, last_logits, probs)) } /// Stateful sampling with the post-chain distribution (#1485): the @@ -505,11 +515,35 @@ pub fn sample_token_with_state_and_distribution( UniquePtr, UniquePtr, ) { + try_sample_token_with_state_and_distribution(logits, config, token_history, state) + .unwrap_or_else(|e| panic!("sampler host read failed: {e}")) +} + +/// Fallible stateful sampling with the post-chain distribution. +pub fn try_sample_token_with_state_and_distribution( + logits: &MlxArray, + config: &SamplingConfig, + token_history: &[i32], + state: &mut Option, +) -> Result { ensure_sampler_state(config, state); let (token, adjusted, probs) = - sample_token_optimized_core_full(logits, config, token_history, state.as_mut(), true); + sample_token_optimized_core_full(logits, config, token_history, state.as_mut(), true)?; let probs = probs.expect("want_distribution returns a distribution on every path"); - (token, adjusted, probs) + Ok((token, adjusted, probs)) +} + +/// Fallible stateful sampling for the engine's per-row chain. +pub fn try_sample_token_optimized_with_state( + logits: &MlxArray, + config: &SamplingConfig, + token_history: &[i32], + state: &mut Option, +) -> Result<(UniquePtr, UniquePtr), String> { + ensure_sampler_state(config, state); + let (token, adjusted, _) = + sample_token_optimized_core_full(logits, config, token_history, state.as_mut(), false)?; + Ok((token, adjusted)) } /// Create the per-sequence [`SamplerState`] on first use when `config` needs @@ -1180,6 +1214,15 @@ fn token_ids_to_host(tokens: &MlxArray) -> Vec { .collect() } +fn try_token_ids_to_host(tokens: &MlxArray) -> Result, String> { + crate::eval_faults::check(crate::eval_faults::FaultPoint::SamplerRead)?; + let bytes = ffi::try_array_to_raw_bytes(tokens).map_err(|e| e.to_string())?; + Ok(bytes + .chunks_exact(4) + .map(|c| i32::from_ne_bytes([c[0], c[1], c[2], c[3]])) + .collect()) +} + /// Batch-safe row-filter hook: apply the enabled RNG-free, history-free /// logit filters to a `[B, V]` (or `[V]`-broadcastable) logit tensor in the /// fixed order pinned by #1375: @@ -1392,7 +1435,7 @@ fn adaptive_p_sample( config: &SamplingConfig, state: Option<&mut SamplerState>, want_distribution: bool, -) -> (UniquePtr, Option>) { +) -> Result<(UniquePtr, Option>), String> { let target = config.effective_adaptive_target(); let decay = config.adaptive_decay.clamp(0.0, 0.99); @@ -1434,15 +1477,16 @@ fn adaptive_p_sample( // Read back the selected token's ORIGINAL probability for the EMA. let idx = ffi::reshape(&ffi::astype(&token, dtype::INT32), &[1, 1]); let sel_p = ffi::take_along_axis(&p, &idx, -1); - ffi::eval(&sel_p); + crate::eval_faults::check(crate::eval_faults::FaultPoint::SamplerRead)?; + crate::try_eval(&sel_p).map_err(|e| e.to_string())?; let orig_p = ffi::item_f32(&sel_p); - let token_id = token_ids_to_host(&token)[0]; + let token_id = try_token_ids_to_host(&token)?[0]; if let Some(a) = adaptive { a.pending = Some((token_id, orig_p)); } let probs = want_distribution.then(|| ffi::softmax(&transformed, -1)); - (token, probs) + Ok((token, probs)) } /// b10621's mirostat bypass path (#1485, upstream `common_sampler_init` with @@ -1471,11 +1515,7 @@ fn mirostat_sample_core( config: &SamplingConfig, state: Option<&mut SamplerState>, want_distribution: bool, -) -> ( - UniquePtr, - UniquePtr, - Option>, -) { +) -> Result { let version = config.effective_mirostat(); let tau = config.mirostat_tau; let eta = config.mirostat_eta; @@ -1492,7 +1532,8 @@ fn mirostat_sample_core( // The plain temperature stage collapses to the argmax; mirostat then // draws the single candidate with observed surprise 0. let token = ffi::fused_sample(&biased, 0.0, 0, 1.0, 0.0); - ffi::eval(&token); + crate::eval_faults::check(crate::eval_faults::FaultPoint::SamplerRead)?; + crate::try_eval(&token).map_err(|e| e.to_string())?; mu -= eta * (0.0 - tau); let probs = want_distribution.then(|| ffi::fused_sample_probs(&biased, 0.0, 0, 1.0, 0.0)); (token, probs) @@ -1519,8 +1560,11 @@ fn mirostat_sample_core( let m = 100.min(vocab); let p = ffi::exp(&logp); let top = ffi::topk(&p, m, -1); - ffi::eval(&top); - let bytes = ffi::array_to_raw_bytes(&ffi::astype(&top, dtype::FLOAT32)); + crate::eval_faults::check(crate::eval_faults::FaultPoint::SamplerRead)?; + crate::try_eval(&top).map_err(|e| e.to_string())?; + crate::eval_faults::check(crate::eval_faults::FaultPoint::SamplerRead)?; + let bytes = ffi::try_array_to_raw_bytes(&ffi::astype(&top, dtype::FLOAT32)) + .map_err(|e| e.to_string())?; let mut top_probs: Vec = bytes .chunks_exact(4) .map(|c| f32::from_ne_bytes([c[0], c[1], c[2], c[3]])) @@ -1567,7 +1611,8 @@ fn mirostat_sample_core( // renormalized distribution. let idx = ffi::reshape(&ffi::astype(&token, dtype::INT32), &[1, 1]); let sel_lp = ffi::take_along_axis(&ffi::log_softmax(&masked, -1), &idx, -1); - ffi::eval(&sel_lp); + crate::eval_faults::check(crate::eval_faults::FaultPoint::SamplerRead)?; + crate::try_eval(&sel_lp).map_err(|e| e.to_string())?; let observed_surprise = -ffi::item_f32(&sel_lp) / LN_2; mu -= eta * (observed_surprise - tau); @@ -1579,7 +1624,7 @@ fn mirostat_sample_core( m.mu = mu; } crate::sampling_dispatch::report_sampling_dispatch(); - (token, biased, probs) + Ok((token, biased, probs)) } /// Top-n-sigma logit filter: keep only the tokens whose logit lies within diff --git a/src/lib/mlxcel-core/src/sampling_row_step.rs b/src/lib/mlxcel-core/src/sampling_row_step.rs index 5cad41b80..948e5d4bc 100644 --- a/src/lib/mlxcel-core/src/sampling_row_step.rs +++ b/src/lib/mlxcel-core/src/sampling_row_step.rs @@ -67,7 +67,8 @@ use crate::ffi; use crate::ffi::MlxArray; use crate::generate::SamplingConfig; use crate::sampling::{ - SamplerState, sample_token_optimized_with_state, sample_token_with_state_and_distribution, + AdaptivePState, MirostatState, SamplerState, try_sample_token_optimized_with_state, + try_sample_token_with_state_and_distribution, }; use cxx::UniquePtr; @@ -109,6 +110,22 @@ impl TokenDraw { } } +/// Failure while preparing a row's draw. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DrawError { + /// The structured-output mask rejected the row. + Mask(String), + /// A sampler host read failed at the backend boundary. + Eval(String), +} + +/// Feedback only: history caches stay in place and synchronize idempotently. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct FeedbackSnapshot { + mirostat: Option, + adaptive: Option, +} + /// One sequence's sampling step and its [`SamplerState`]. /// /// Used by: `engine::sample_and_finish_row` @@ -151,11 +168,26 @@ impl RowSampler { self.state.as_ref() } + pub(crate) fn feedback_snapshot(&self) -> FeedbackSnapshot { + FeedbackSnapshot { + mirostat: self.state.as_ref().and_then(|s| s.mirostat), + adaptive: self.state.as_ref().and_then(|s| s.adaptive.clone()), + } + } + + pub(crate) fn restore_feedback(&mut self, snapshot: FeedbackSnapshot) { + if let Some(state) = self.state.as_mut() { + state.mirostat = snapshot.mirostat; + state.adaptive = snapshot.adaptive; + } + } + /// Draw one token from `logits` (one row) under `config`. /// /// `history` must already end with the previous token (see /// [`Self::needs_host_token_before_sample`]). `mask` runs before the - /// chain; its error is returned unchanged and nothing is sampled. + /// chain; its error is returned as [`DrawError::Mask`] and nothing is + /// sampled. A backend read failure is [`DrawError::Eval`]. /// `want_distribution` attaches the post-chain distribution, drawn in the /// same chain pass so both share one XTC gate. /// @@ -169,13 +201,18 @@ impl RowSampler { history: &[i32], mask: Option<&mut dyn LogitMask>, want_distribution: bool, - ) -> Result { + ) -> Result { let Some(mask) = mask else { - return Ok(self.sample(logits, config, history, want_distribution)); + return self + .sample(logits, config, history, want_distribution) + .map_err(DrawError::Eval); }; let vocab_size = ffi::array_shape(logits).last().copied().unwrap_or(0) as usize; - let masked = mask.apply(ffi::copy(logits), vocab_size)?; - Ok(self.sample(&masked, config, history, want_distribution)) + let masked = mask + .apply(ffi::copy(logits), vocab_size) + .map_err(DrawError::Mask)?; + self.sample(&masked, config, history, want_distribution) + .map_err(DrawError::Eval) } /// [`Self::draw`] without a mask or a distribution, confirmed at once as @@ -189,10 +226,10 @@ impl RowSampler { logits: &MlxArray, config: &SamplingConfig, history: &[i32], - ) -> TokenDraw { - let draw = self.sample(logits, config, history, false); + ) -> Result { + let draw = self.sample(logits, config, history, false)?; self.accept_drawn(); - draw + Ok(draw) } /// Apply `token_override` to the host value of the last draw and confirm @@ -224,26 +261,31 @@ impl RowSampler { config: &SamplingConfig, history: &[i32], want_distribution: bool, - ) -> TokenDraw { + ) -> Result { if want_distribution { let (token, adjusted_logits, distribution) = - sample_token_with_state_and_distribution(logits, config, history, &mut self.state); - TokenDraw { + try_sample_token_with_state_and_distribution( + logits, + config, + history, + &mut self.state, + )?; + Ok(TokenDraw { token, adjusted_logits, distribution: Some(distribution), - } + }) } else { // With no state this is `sample_token_optimized` exactly: the // chain's lazy creator allocates nothing for a config that needs // no state. let (token, adjusted_logits) = - sample_token_optimized_with_state(logits, config, history, &mut self.state); - TokenDraw { + try_sample_token_optimized_with_state(logits, config, history, &mut self.state)?; + Ok(TokenDraw { token, adjusted_logits, distribution: None, - } + }) } } diff --git a/src/lib/mlxcel-core/src/sampling_row_step_tests.rs b/src/lib/mlxcel-core/src/sampling_row_step_tests.rs index 70210830a..b9d57c24c 100644 --- a/src/lib/mlxcel-core/src/sampling_row_step_tests.rs +++ b/src/lib/mlxcel-core/src/sampling_row_step_tests.rs @@ -401,5 +401,5 @@ fn the_mask_runs_before_the_chain_and_its_error_aborts_the_draw() { .draw(&logits, &greedy, &[], Some(&mut refuse), false) .err() .expect("a refusing mask aborts the draw"); - assert_eq!(err, "no token is allowed"); + assert_eq!(err, super::DrawError::Mask("no token is allowed".into())); } diff --git a/src/server/batch/scheduler/decode_tick.rs b/src/server/batch/scheduler/decode_tick.rs index 1581555f1..01e4222bb 100644 --- a/src/server/batch/scheduler/decode_tick.rs +++ b/src/server/batch/scheduler/decode_tick.rs @@ -611,7 +611,13 @@ impl BatchScheduler { let next = self.prime_lookahead_with_input(seq_ids, &next_input, params); // Step 2: read step n's tokens to host (the sync point) and finish-check. - let toks = mlxcel_core::engine::tokens_to_host(&la.tokens); + let toks = match mlxcel_core::engine::try_tokens_to_host(&la.tokens) { + Ok(tokens) => tokens, + Err(msg) => { + let _ = self.record_eval_outcome(Err(msg)); + Vec::new() + } + }; let mut finishing = toks.len() != seq_ids.len(); if !finishing { for (i, &seq_id) in seq_ids.iter().enumerate() {