Repository navigation
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
Description
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 tagnvfp4._KIND_KERNELkeeps the first entry (:40-42), sokind_kernel_for("nvfp4")is alwaystriton.python/freetoken/engine/engine.py:1596-1619_adjust_ftw_quant_backend: underautoit 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-272and:285-287write the alphas of the bound kernel (w4a4on sm_120 underauto), andcheckpoint/ftw.py:696-701passes them toExpertBanks. Thetritonlayout declares no alpha roles, somoe/expert_banks.py:488-496never 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 annvfp4FTW as compatible with bothtritonandw4a4. Underauto, restrict the table to those two and letselect_kernelchoose withworth_it()(sow4a4on sm_120). Alloww4a4only when the index hasgate_up_alpha/down_alphaentries. Without them, refusew4a4with 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
w4a4its own tag (for examplenvfp4_w4a4) that reads back asw4a4on GPUs where it is usable and astritonelsewhere.
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