Module diagnostics 

Optional CPU-reference comparisons, research dumps and per-step traces.

Functions

argmax 🔒

dump_decode 🔒

dump_prefill_chunk 🔒

Record prompt argmax rows and the final long-context attention selection.

dump_run 🔒

logits_diff 🔒

qwen4_exp_check_rows 🔒

Compare the GPU’s rows against the CPU reference, which forwards the same tokens one at a time. With next (the token after each row), the CPU also runs the MTP head per row and the last row’s draft logits are compared with the GPU’s.