Skip to content

An FTW checkpoint can never serve the W4A4 NVFP4 MoE prefill, and the refusal tells you to reconvert with the option that produced it #766

Open

Description

@gdevenyi

What happened

On sm_120 the automatic kernel choice for NVFP4 experts with static input scales is w4a4 (#736). ft checkpoint with the automatic choice therefore binds w4a4. It writes the nvfp4 bank format plus the gate_up_alpha / down_alpha entries. When that FTW is served, nvfp4 maps back to triton, so the W4A16 prefill runs. If you ask for --quant-backend moe.nvfp4=w4a4, the server refuses to start, and the error tells you to reconvert with ft checkpoint --quant-backend moe.nvfp4=w4a4. That reconversion writes nvfp4 again and fails with the same error. Nothing you can pass to either command gives an FTW that runs the W4A4 prefill. The alpha entries load (they are small) but no kernel reads them.

#736's "Known issues, not addressed in this PR" comment notes the W4A16 fallback ("FTW checkpoints keep the W4A16 prefill"). This issue tracks that fallback, together with the error message that cannot be followed.

Minimal reproduction (CPU only, PYTHONPATH=python):

import freetoken.checkpoint.ftw as ftw
from freetoken.engine.engine import _adjust_ftw_quant_backend
from freetoken.layers.quantization import QuantKind
from freetoken.layers.quantization.quant_backend import QuantBackend
from freetoken.moe.legacy_format import kind_kernel_for, legacy_format_for

fmt = legacy_format_for(QuantKind.NVFP4, "w4a4")
print("ft checkpoint, kernel w4a4 -> quant_format", repr(fmt), "-> served as", kind_kernel_for(fmt))
ftw.ftw_quant_format = lambda path: fmt  # an FTW converted with moe.nvfp4=w4a4 (or auto on sm_120)
print("serve, auto ->", _adjust_ftw_quant_backend("/ftw", QuantBackend.parse(None)))
try:
    _adjust_ftw_quant_backend("/ftw", QuantBackend.parse("moe.nvfp4=w4a4"))
except ValueError as e:
    print("serve, moe.nvfp4=w4a4 -> ValueError:", e)
ft checkpoint, kernel w4a4 -> quant_format 'nvfp4' -> served as (<QuantKind.NVFP4: 'nvfp4'>, 'triton')
serve, auto -> QuantBackend(items=(((<LayerKind.MOE: 'moe'>, <QuantKind.NVFP4: 'nvfp4'>), 'triton'),))
serve, moe.nvfp4=w4a4 -> ValueError: the FTW checkpoint's expert banks were packed for nvfp4 / triton but --quant-backend asks for w4a4; drop the entry, or reconvert with ft checkpoint --quant-backend moe.nvfp4=w4a4

Cause (current main, cc56b6b)

  • python/freetoken/moe/legacy_format.py:32-34: (NVFP4, "triton") and (NVFP4, "w4a4") share the tag nvfp4. _KIND_KERNEL keeps the first entry (:40-42), so kind_kernel_for("nvfp4") is always triton.
  • python/freetoken/engine/engine.py:1596-1619 _adjust_ftw_quant_backend: under auto it pins MoE/NVFP4 to that kernel. Any other request raises, and the message suggests a reconversion that cannot give a different tag.
  • python/freetoken/checkpoint/convert.py:268-272 and :285-287 write the alphas of the bound kernel (w4a4 on sm_120 under auto), and checkpoint/ftw.py:696-701 passes them to ExpertBanks. The triton layout declares no alpha roles, so moe/expert_banks.py:488-496 never attaches them to a layer.

Proposed fix

The banks are identical, so the choice between w4a4 and triton does not have to be fixed at conversion time:

  • In _adjust_ftw_quant_backend, treat an nvfp4 FTW as compatible with both triton and w4a4. Under auto, restrict the table to those two and let select_kernel choose with worth_it() (so w4a4 on sm_120). Allow w4a4 only when the index has gate_up_alpha / down_alpha entries. Without them, refuse w4a4 with a message that says the FTW was converted without the input scales (and how to get them), not one that suggests the same command.
  • Or give w4a4 its own tag (for example nvfp4_w4a4) that reads back as w4a4 on GPUs where it is usable and as triton elsewhere.

A related option: ft checkpoint could store the input scales whenever the checkpoint has them, so that an FTW converted on a GPU that is not sm_120 also gets W4A4 when it is served on sm_120.

Environment

Seen serving nvidia/Qwen3.8-Flash-Next-NVFP4 on an RTX 5080 16 GB (sm_120, driver 615.78.08), torch 2.14.1+cu130; reproduction CPU-only on main cc56b6b.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions