Skip to content

Check penalty oracle decisions against the actual target forward logits - #526

Open
linson007 wants to merge 4 commits into
incoai:mainfrom
linson007:deepinfer-penalty-oracle
Open

linson007 wants to merge 4 commits into
incoai:mainfrom
linson007:deepinfer-penalty-oracle

Conversation

@linson007

@linson007 linson007 commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Commit attribution was corrected to linson007 with a metadata-only history rewrite. File trees are unchanged; historical benchmark source hashes and CI links below retain their original recorded scope.

The real-model penalty oracle can reject a correct incremental selection because it recomputes its reference logits through cold prefill. On current main 1482f6bb, the M5 Pro/Qwen3.6 GGUF BF16-KV/FP32-state run fails with cold-score loss 0.205502 > 0.1, although an independent reference over that forward's actual logits agrees with every checked penalty result and greedy winner.

This PR replaces that hard cross-graph assertion with a same-forward policy check and keeps cold-score disagreement as a diagnostic. It does not increase the old tolerance. It changes test-only instrumentation and the oracle, with no production runtime, shader or default change.

Refs #437; discovered during #508/#523 validation. #437's Ornith/M5 Max reproduction remains a separate configuration; this PR records a new shader-validated Qwen/M5 Pro reproduction rather than claiming to have rerun that reporter's model.

Check flow

flowchart LR
    A["Actual target forward"] --> B["Raw-logit snapshot"]
    A --> C["GPU penalty policy"]
    C --> D["Penalized logits and greedy choice"]
    B --> E["Independent CPU policy using fixture counts"]
    E --> F["Check same-row arithmetic and winning token"]
    D --> F
    G["Cold-prefill recomputation"] --> H["Retained drift diagnostic"]
Loading

The snapshot belongs to the instrumented test build. CPU counts come from the
fixture's prompt, emitted transcript, pending anchor and accepted proposal
prefix; they are not read back from the runtime's penalty table.

Review the code

  1. RuntimeInstrumentation.hpp / instrumented Runtime.cpp: an explicitly installed observer snapshots the forward's raw target logits before penalties. Captured Metal buffer views remain owned through the host callback. The oracle copies values into its own storage. All fields, calls and definitions are behind SPLASH_RUNTIME_INSTRUMENTATION; an unset observer adds no host step.
  2. dev/PenaltyOracle.hpp: independent fp64 arithmetic over raw logits, prompt membership and generated-token counts; compare each vocabulary logit with an eight-float-epsilon allowance for fp32 arithmetic and require the selected token to attain the actual penalized maximum.
  3. model_runtime_oracle_test.mm: rebuild expected counts from the fixture transcript, pending anchor and accepted proposal prefix, including initial selection and all three preemption/replay modes. Check initial and retained verify rows, not stale/unexecuted suffixes. The original cross-graph loss is still printed.
  4. penalty_oracle_test.cpp / dev/native.mk: hand-calculated positive/negative fixtures and rejection tests for missing repetition, stale counts/prefixes even with an unchanged winner, wrong selections, nonfinite corruption and excess arithmetic error. The unit joins the CPU suite; oracle/header dependencies ensure reference edits rebuild the executable.

The observer adds a host step, so these are correctness runs, not timing measurements. The old cold-score loss still reproduces in the instrumented run, demonstrating that the new check is not passing merely because the prior disagreement disappeared.

Reproduction and evidence

M5 Pro, 20 GPU cores, 48 GiB; macOS 26.6.0, clang 21.0.0, Metal 32023.921, SDK 27.0, deployment target 26.4.

Target unsloth/Qwen3.6-35B-A3B-GGUF:UD-IQ3_XXS, revision a483e9e6cbd595906af30beda3187c2663a1118c; draft incoai/Qwen3.6-35B-A3B-DFlash2, revision 51ef7b6923ad6c14cb1bb41c37a9041446496ab6.

make -j4 build/splash build/engine-tests/model-runtime-oracle \
  build/engine-tests/penalty-oracle build/engine-tests/target-sampling
make test-engine-cpu
MTL_SHADER_VALIDATION=1 build/engine-tests/target-sampling build/splash.metallib
MTL_SHADER_VALIDATION=1 build/engine-tests/model-runtime-oracle \
  build/splash.metallib "$MODEL_ROOT" \
  --kv-format bf16 --state-format fp32 --ane-ffn-share 0 --max-context 32768

Initial model fix 17966d05: PASS for the complete formerly failing invocation. Twelve penalty/preemption runs check 248 rows, every vocabulary logit in each, with maximum normalized arithmetic error 5.90884e-08. The cold-score diagnostic remains 0.205502. Original failing and corrected full logs, model manifest, commands, source/binary hashes and production-preprocessor parity are in dev/benchmarks/results/qwen-penalty-oracle-m5-pro/.

The final reference edge correction c20ea1e1 leaves prompt-only logits untouched when repetition is inactive, including positive infinity. Updated CPU unit and ASAN/UBSAN pass. Four representative final-source model cells now pass, after #523's current-main serving checks. The earlier model run remains explicitly scoped to its own source. Full final logs, commands, all four model-manifest identities and binary hashes are attached in evidence commit d34b99fc.

Validation and limits

  • Full local CPU suite: pass; final updated helper unit: pass.
  • Helper ASAN/UBSAN: pass, command/log attached.
  • Shader-validated target-sampling: pass.
  • Initial complete real-model correction run: pass as detailed above.
  • Production Runtime.cpp preprocessed tokens are identical to base with instrumentation unset (blank-line/line-marker normalization only): SHA256 931abfa5be9f190a387fe30460a4e498f5ac73dedd18b000e78f00a4d63e8ec6, 4,279,706 bytes.
  • Final code c20ea1e1: four complete representative model oracles PASS. MoE GGUF BF16-KV/FP32-state: 248 strict rows; MoE MLX INT8-KV/BF16-state: 294; dense MLX BF16/BF16: 295; dense GGUF INT8/FP32: 300. Total 1,137 rows, maximum normalized arithmetic error 5.94946e-08. Each cell covers all twelve penalty/preemption cases. These are four representative combinations, not all sixteen KV/state combinations rerun.
  • Final head d34b99fc passes all four CI jobs: production/native/CPU/sanitizers and Python 3.12/3.13/3.14 (run 39602804636). The hardware release gate is skipped, not counted as a pass.
  • Other machines and the original Ornith report remain unexecuted locally.

This establishes penalty arithmetic/greedy decisions for the actual forward. It does not locate or repair the underlying prefill/verify numerical drift, assert sampling distributions or model quality, or qualify performance.

@linson007

Copy link
Copy Markdown
Contributor Author

Rechecked against main 1482f6bbc2abf12375332725302019594150b6cb and the maintainer's #280/#283 closure.

This branch already contains current main: 4 commits ahead, 0 behind. Its source is unchanged; all four CI jobs remain green on d34b99fc. The existing four-model correction evidence therefore retains its original source identity and is not relabelled as a new run.

#487's adaptive MoE row policy is independent of this penalty-oracle correction. The observer reads the command's actual logical verify-row count and checks initial/retained rows; it does not reinstate the old fixed-threshold feature. Production Runtime.cpp remains unchanged after preprocessing with instrumentation unset. The original fifteen-pass/one-failed upstream matrix remains recorded separately from the four passing corrected representative cells.

@linson007
linson007 force-pushed the deepinfer-penalty-oracle branch from d34b99f to ad3e4f1 Compare October 21, 2026 18:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant