Skip to content

Fix PPL VLM last-position scoring: rewind-free incremental teacher forcing - #19

Open
amd-genmingz wants to merge 1 commit into
AMDmoore:mainfrom
amd-genmingz:fix/ppl-vlm-rewind-free-teacher-forcing
Open

Fix PPL VLM last-position scoring: rewind-free incremental teacher forcing#19
amd-genmingz wants to merge 1 commit into
AMDmoore:mainfrom
amd-genmingz:fix/ppl-vlm-rewind-free-teacher-forcing

Conversation

@amd-genmingz

Copy link
Copy Markdown

Summary

  • Rewrites the last-position-only branch of compute_sample_nll in
    tests/PPL/perplexity_vlm.py to score targets without rewind_to.
  • Uses a single generator with incremental teacher forcing: prefill the
    prompt once, then append_tokens the true previous target token per
    step and read the final-position logits. generate_next_token is never
    called, so no sampled token is inserted into the sequence.

Why

The previous last-position path relied on rewind_to + append_tokens to
recover per-position logits. On gemma4 multimodal + amdgpu EP with
past_present_share_buffer=false, the OGA per-layer KV-cache branch does
not maintain the sequence-length bookkeeping (shape_[2]) across a rewind.
After a rewind, OGA's pre-allocated present.* shape disagrees with the
EP-computed shape (e.g. {1,8,277,256} vs {1,8,296,256}), and the native
layer fatally aborts with 0xC0000409 (STATUS_STACK_BUFFER_OVERRUN).

What changed

  • Drop all rewind_to calls from the last-position branch.
  • Prompt prefill (and the vision encoder / image features) now runs once;
    each target position costs a single incremental decode step. Cost goes from
    O(target_len) full prefills back to O(1) prefill + O(target_len) cheap
    single-token decodes.
  • The all-position (if) branch is unchanged; models that emit full
    [B, T, V] logits still take the single-pass fast path.
  • Dump fields (target_ids / target_nll / top1_ids / topk) are
    unchanged, so downstream compare_vlm_ep_cpu.py is unaffected.
  • Minor: fix argparse help string (%% escaping and en-dash normalization).

Correctness

  • After set_inputs(prompt), final-position logits = P(? | prompt) =
    prediction for target[0].
  • After append_tokens([target[j-1]]), final-position logits =
    P(? | prompt + target[:j]) = prediction for target[j].
  • Mathematically equivalent to the original rewind logic and to the
    all-position path (same set of target-position NLLs).

The last-position-only branch in compute_sample_nll scored PPL targets
with rewind_to + append_tokens. On gemma4 multimodal + amdgpu EP with
past_present_share_buffer=false, the per-layer KV-cache branch does not
maintain the sequence-length bookkeeping across a rewind, so the
pre-allocated present.* tensor shape disagrees with the EP-computed
shape and the native layer fatally aborts (0xC0000409).
Replace it with a single generator doing incremental teacher forcing:
prefill the prompt once (vision encoder + image features computed a
single time), then append_tokens the true previous target token per
step and read the final-position logits. Only append_tokens is used
(never generate_next_token), so no sampled token is inserted and
rewind_to is never called. This keeps the expensive prefill at O(1)
instead of O(target_len) and is mathematically equivalent to the
all-position path.
Also fix an argparse help string (%% escaping and en-dash) in the
CLI section.
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