Skip to content

Reduce unused vocabulary projections during Qwen prefill (~8% speedup) - #614

Open
aleroot wants to merge 1 commit into
ml-explore:mainfrom
aleroot:qwen_prefill_voc
Open

aleroot wants to merge 1 commit into
ml-explore:mainfrom
aleroot:qwen_prefill_voc

Conversation

@aleroot

@aleroot aleroot commented Sep 7, 2026 •

Copy link
Copy Markdown
Contributor

Proposed changes

Qwen currently computes output scores for prompt positions that normal generation discards. This change reduces that work, making prompt processing faster and reducing peak memory usage.

Applies to supported quantized Qwen3 and Qwen3.5 models, including Qwen3.5’s vision-language path. It keeps the same calculation order for the final scores and preserves full outputs where needed, including speculative decoding.

On my Apple M2 with 16GB RAM, using 512-token text prompts and alternating baseline/optimized runs:

Model Before → after Lower prefill latency
Qwen3-4B-4bit 3.45 → 3.17 s 8.1%
Qwen3.5-9B-4bit, text model 5.42 → 4.84 s 10.7%
Qwen3.5-9B-4bit, vision-language model with text input 5.40 → 4.82 s 10.7%

Peak memory decreased by approximately 130, 185, and 212 MiB respectively. Gains depend on prompt length and prefill chunking.

Exact comparisons passed for final scores, cache contents, and generated tokens on both cached checkpoints.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

AI usage

  • I have read this PR description in full and approve it as my own, and it
    accurately describes the code changes.
  • AI usage disclosure:

Keep aligned final projection tiles for Qwen3 and Qwen3.5 generation,
preserving full-output paths and cache updates. Add exact parity,
fallback, and local checkpoint tests.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants