Skip to content

feat: NVIDIA Qwen3.8 Flash-Next MTP loader - #553

Open
jkyamog wants to merge 1 commit into
1CatAI:mainfrom
jkyamog:fix/nvidia-flash-next-sm70
Open

jkyamog wants to merge 1 commit into
1CatAI:mainfrom
jkyamog:fix/nvidia-flash-next-sm70

Conversation

@jkyamog

@jkyamog jkyamog commented Sep 7, 2026

Copy link
Copy Markdown

Purpose

Fix the official NVIDIA Qwen3.8-Flash-Next checkpoint on SM70 when MTP is
enabled. Its MTP routed experts are FP8 E4M3 block-scaled weights with 128-wide
blocks, while TP4 shards the 640-wide expert projections into 160-wide slices.
The existing mixed-quantization path missed the runtime MTP layer metadata and
could not dispatch this block-scale format, so MTP K=4 failed during loading.

The change:

  • remaps checkpoint mtp.layers.0 metadata to the runtime MTP layer index;
  • converts only the official MTP block-FP8 expert weights to resident FP16 at
    load time, applying the checkpoint scales once before TP sharding;
  • leaves the main NVFP4 model path and non-SM70 behavior unchanged; and
  • adds focused validation for ordering, scale expansion, TP slicing, and
    malformed/incomplete tensor pairs.

This is the correctness-first FP16 path. Native block-FP8 dispatch remains a
future optimization requiring block-aligned sharding and separate accuracy and
performance qualification.

Test Plan

Run from the repository root in the CUDA/SM70 test container:

pytest -q tests/quantization/test_modelopt_mtp_block_fp8.py
ruff check vllm/models/qwen4_exp/nvidia/mtp.py \
  vllm/models/qwen4_exp/nvidia/mtp_fp8.py \
  tests/quantization/test_modelopt_mtp_block_fp8.py
git diff --check origin/main...HEAD

Build and runtime qualification:

TORCH_CUDA_ARCH_LIST=7.0 MAX_JOBS=4 docker build ...
./vllm/vllm_ctl.sh status
node vllm/tmp/probe_pi_stream.mjs qwen-3.8-flash

The live NVIDIA checkpoint gate should show TP4, MTP K=4, the resident-FP16
loader witness, health 200, and zero restarts. Re-run the standard
llama-benchy stan command at the selected power limit for performance data;
the archived 300 W samples are a reference, not a claim about the current
250 W run.

Test Result

Previously validated in the combined image based on the earlier upstream
base:

  • 15 focused NVIDIA quantization tests passed.
  • A real NVIDIA checkpoint probe matched an independent FP32-scale-then-FP16
    reference bit-for-bit for 27 sampled tensors and all 108 TP4 slices; outputs
    were finite.
  • The four-V100 service loaded all 11 shards, selected the resident-FP16 MTP
    path, completed exact-shape MTP4 warmup and CUDA graph capture, returned
    health 200, and had zero restarts/OOMs.
  • The standard pure stream-timestamp stan measurement at 300 W averaged
    3347.5 PP tok/s, 135.3 TG tok/s, and 20.25 s TTFR/E2E over three samples.

@github-actions

github-actions Bot commented Sep 7, 2026

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

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