Implement Multi Token Prediction

[tomiock]

use the already existing MTP module of the model to implement the loss during training via next token prediction

• implement in qwen3-vl and qwen3.5
• verify decoding speedup
• test quick training run

[tomiock]

Quick benchmark test:

(torch11) tockier@cudahpc51:~/vlm_training_upstream$ CUDA_VISIBLE_DEVICES=7 python eval/benchmark_mtp_decode.py \
  --model_dir /data/151-1/users/tockier/qwen_finetune/cache/qwen35_2b \
  --model_type qwen3_5 --num_tokens 128 --out eval/mtp_decode_report.md
model=qwen3_5  mtp_num_hidden_layers=1

======================================================================
  MTP Self-Speculative Decoding — qwen3_5
======================================================================
  greedy · no KV-cache (eager) · bfloat16 · 128 tokens/prompt

  ▶ 1.63× faster decoding  (33.4 → 54.6 tok/s)  ·  lossless (identical output)
  ▶ draft acceptance 77%  ·  1.74 tokens per model forward (vs 1.00 baseline)

  Sample generations (prompt → continuation; baseline and MTP produce the same text)
  ------------------------------------------------------------------
  [1] The history of the Roman Empire began when
      the Roman Republic was founded in 509 BC. The Roman Republic was a political system that existed from 509 BC to 27 BC. The Roman Empire was a political system that existed from 27 BC to 476 AD. The Roman Empire was a political system that existed from 476 AD to 1453 AD. The Roman Empire was a political system that exis …
      1.73× · accept 86% · ✓ identical
  [2] Here is a simple recipe for chocolate chip cookies. First,
      you need to mix the ingredients. Then, you need to bake them. Finally, you need to eat them. The ingredients are: - 1 cup of flour - 1 cup of sugar - 1 cup of butter - 1 cup of chocolate chips - 1 cup of eggs - 1 cup of milk - 1 cup of vanilla extract - 1 cup of salt - 1 cup of baking soda - 1 cup of baking powder - 1 …
      1.67× · accept 80% · ✓ identical
  [3] In the field of machine learning, a transformer is
      a type of neural network architecture that has become increasingly popular in recent years. It is known for its ability to process and analyze large amounts of data, making it a powerful tool for various tasks such as natural language processing, computer vision, and speech recognition.  One of the key components of a …
      1.66× · accept 82% · ✓ identical
  [4] Once upon a time, in a small village by the sea, there lived
      a curious little boy named Tom. Tom loved to explore the world around him, especially the things that happened near the water. One sunny afternoon, while playing in the sand, he noticed something strange.  Tom saw a small boat floating on the water. He wondered, "Why is the boat floating?" He knew that the water was ve …
      1.50× · accept 62% · ✓ identical

  Per-prompt metrics
  prompt                              base t/s  spec t/s  speedup  accept
  ------------------------------------------------------------------
  The history of the Roman Empire         33.1      57.3    1.73×     86%
  Here is a simple recipe for choc        33.5      56.1    1.67×     80%
  In the field of machine learning        33.7      55.7    1.66×     82%
  Once upon a time, in a small vil        33.2      49.8    1.50×     62%
  OVERALL                                 33.4      54.6    1.63×     77%
======================================================================