refactor(quant): config, scheme and method layers for quantization - #418
Merged
Merged
Conversation
…tization QuantConfig parses the checkpoint's quantization_config into a QuantScheme per projection; every Linear*, ParallelLMHead and MoE layer takes quant_config and a prefix, declares its tensors through a QuantMethod and selects its kernel at construction. MoE kernels own the expert bank layout, pack and apply, so the offload cache, the FTW converter and the CPU executor follow one layout. - --quant-backend layer[.kind]=name selects kernels (--nvfp4-backend forwards to it) - --moe-strategy replaces --moe-backend; the old spelling stays as an alias - --moe-cpu-layers auto for Windows/WSL pin budgets; --moe-cache-auto respects the kernel's slot limit - removes the per-family quantized layer classes, bank providers, format env switches and GLM runtime fp8 - verified with AIME25 pass@3 on the models.md checkpoints and greedy token parity with main
Collaborator
|
@jason-fxz Huge PR but great abstraction! A round of validation is needed. |
- bias is an architecture property and is never quantized: Linear and LM head allocate it themselves after create_weights, the five Linear methods only read layer.bias in apply - MoELayer / MoEConfig gain has_bias (gpt-oss passes True); the mxfp4 kernels and method read it instead of a scheme role - the bias role leaves every scheme builder; roles now name checkpoint tensors only - _common_reject refuses a biased layer for every kernel without a bias epilogue, so only triton_gptoss accepts one
- the kernel calls torch._scaled_mm, which torch routes to cuBLASLt for tensorwise scales and to its own CUTLASS kernel for rowwise scales, so cublaslt only named half of it - follows vLLM, which names these kernels after the library that ships the op, and matches the torch kernel of linear.none - --quant-backend linear.fp8_tensor=torch
This was referenced Sep 9, 2026
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 10, 2026
Brings the quantization refactor (FlashML-org#418 config/scheme/method layers, FlashML-org#427 the QuantConfig reaching the weight readers) and FlashML-org#426, which takes qwen4_exp's expert quant kind from the checkpoint QuantConfig. Resolutions, all taken from the rebased PR branches so the deploy tree and the upstream PRs stay identical where they overlap (rb/tp, rb/vision, rb/fp8 -> rb/all): - qwen4_exp attention/gdn/moe/config/model, layers/linear.py, layers/embedding.py, models/nvfp4_banks.py: the rebased versions. Notably o_proj stays LinearOProj (row-parallel); main's LinearReplicated is correct only at TP=1. - NVFP4 expert TP sharding moved to the new seam: the piece stream is sliced along the intermediate axis in nvfp4_banks, and the kernel sizes its banks from MoEConfig.local_intermediate. The old _Placer/_alloc_nvfp4_host_banks are gone with the functions they lived in. - models/quant_linear.py: deleted, as upstream did; nothing imports it. - DROPPED, superseded upstream: the FlashML-org#320 cherry-pick (_dense_is_block_fp8, _block_fp8_fusions, _load_maybe_block_fp8). FlashML-org#320 is CLOSED upstream because FlashML-org#426 does the same job through QuantConfig. This checkpoint's dense side is bf16 (a 292-entry modelopt ignore list), so none of those paths were exercised here. - KEPT: the --num-tokens/--num-pages KV reserve fix (issue FlashML-org#383, still unfixed on main), FlashML-org#231's collect_decode_freq, FlashML-org#169's prefill_warmup, and every other deploy-only commit -- they merged clean. - tests/engine/test_cache_budget.py keeps both the KV-reserve test and upstream's new slot_limit test; _resolve_auto_moe_cache_size's new method arg is optional. deploy/chatdnp stays at adc32da as the rollback point. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
This was referenced Sep 10, 2026
This was referenced Sep 10, 2026
This was referenced Sep 10, 2026
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 10, 2026
…e dropped layers/quantization/moe/nvfp4.py is new in upstream main (FlashML-org#418), so it did not exist on the deploy branch's side of the merge and git took upstream's copy without a conflict -- silently discarding the rb/tp change that sizes the banks from MoEConfig.local_intermediate. The result was a tree that slices the expert piece stream per rank (nvfp4_banks._tp_shard, which DID come across) while still allocating full-width banks: at TP=2 local_intermediate=320 but gate_up was (1280, 1280) rather than (640, 1280). Checked the whole class of error rather than this one file: of the 42 files rb/all touches, this is the only one that is new in main, and every other difference against the deploy branch is deploy-only work rb/all does not carry (verified e.g. distributed_timeout=1800 is present). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
CraigStone-Dev
pushed a commit
to CraigStone-Dev/FreeToken
that referenced
this pull request
Sep 11, 2026
Re-land of PR FlashML-org#337 onto the post-FlashML-org#418 quantization refactor: the disk tier serves experts that do not fit in pinned RAM -- the RAM bank holds the first --expert-ram-experts per layer (pinned), the rest stay in the original safetensors checkpoint and are fetched O_DIRECT -> pinned staging -> H2D into the slot the LRU kernel already assigned, shrinking the miss list so the existing PCIe copy_missing path only moves the RAM-resident misses. Adapted to the new architecture: * Bank layout is owned by the expert kernel (BankSpec per role); the tier speaks the native NVFP4 (triton) layout, which is unchanged by the refactor. The index and fetch path are written against it. * The per-family nvfp4_expert_source_spec hook from round 4 is subsumed by upstream's own nvfp4_expert_spec hook (added in the refactor); the index resolves the spec through the same hook the reader uses, so index and loader read the same rows. Built in _method_expert_banks next to the release, so a family that releases rows without an index fails loudly instead of serving zeroed experts. * Load-time RAM peak is the K/E prefix, not the full expert set: the NVFP4 reader never reads rows [K, E) (serial: no get_tensor; parallel: filtered at the reader) and yields an empty piece per skipped expert so build_expert_banks completes the layer, pins only the prefix (PinPipeline(prefix_rows=K)) and releases the tail (release_bank_tails). * Host banks are MAP_PRIVATE|MAP_ANONYMOUS (MT-z round 3): release_range is one madvise(MADV_DONTNEED) that really frees; the mincore startup check (check_tail_unbacked) warns if anything backed the released tail. * FT_DISK_TIER_VERIFY [copy-miss] probe gated on the tier + not capturing (round 3 finding); --moe-disk-tier preconditions collected into one error. CPU suite: test_disk_tier.py 8/8, test_disk_tier_families.py 15/15, test_offload.py green (incl. the new probe-gate regression test).
tomasuz
pushed a commit
to tomasuz/FreeToken
that referenced
this pull request
Sep 11, 2026
Upstream FlashML-org#418/FlashML-org#427/FlashML-org#438 moved expert quantization into MoEMethod/MoEKernel (layers/quantization/moe): bf16, NVFP4 and fp8 experts are packed by build_expert_banks and applied through quant_method.apply; only GGUF still loads through a format provider. --moe-backend became --moe-strategy and EngineConfig.moe_backend is folded into moe_strategy and left None. Resolution keeps every feature of this branch on top of that design: - engine: take upstream's shared_offload_method / _check_pin_budget / PinFailed / layout+max_slots path; keep the resident tier, worker executors, in-place worker layers, shared banks and decode-frequency stats. load_expert_banks is called once, inside shared_banks(), with both method= and resident_layers=. Our helpers read moe_strategy. - expert_banks: upstream build_expert_banks plus ResidentUploader; the resident upload rides host_banks' PinPipeline sink, which upstream's builder already drives. The bf16/nvfp4/ds_fp4 providers are gone as upstream intended; one GGUF provider serves every native ggml quant. - layers/moe: the format-tag dispatch keeps only the GGUF branch; the split/worker decode path already goes through _expert_gemm and so through quant_method.apply. - models/weight: load_gguf_moe_expert_sources, with upstream's load_q4_0_moe_expert_sources kept as an alias. - moe/fused_q4_0: add upstream's fused_experts_gguf_q4_0 entry point. - moe/_worker_main: take NVFP4 banks by kernel role (gate_up, down, ...), falling back to the FTW names. - tests: fake configs in test_resident / test_worker_placement use moe_strategy. Checked on tm (ROCm 7.2, torch 2.11): no undefined names in the 180 changed files; tests/engine + tests/moe show no failure that is not also present on the pre-merge branch or on pure upstream main in the same environment (51 failed / 285 passed; test_nvfp4_moe_vec.py and test_warp_width.py excluded -- both segfault identically before the merge). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014Tt8JH1wvkbStVdTSzHNPE
This was referenced Sep 12, 2026
nomanoma121
pushed a commit
to nomanoma121/My-FreeToken
that referenced
this pull request
Sep 16, 2026
Upstream FlashML-org#418 moved kernel selection into a QuantConfig / QuantScheme / QuantMethod layer and deleted the per-family quantized layer classes --dense-quant was overriding. The flag now wraps the checkpoint's quant config instead (layers/quantization/at_load.py): it reports per-row fp8 for the projections a checkpoint left bf16, and the loader converts them as before. The fork's fp8 and NVFP4 head classes give way to ParallelLMHead's quant method, and the pipeline's expert-bank window moved from each family's reader to expert_banks._local_pieces. Upstream also stopped picking CPU MoE layers on its own: pass --moe-cpu-layers auto where the banks exceed the host's pin budget, or the boot stops and asks for it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
nomanoma121
pushed a commit
to nomanoma121/My-FreeToken
that referenced
this pull request
Sep 16, 2026
The fork already fixed this once (12.50 -> 14.59 tok/s on gpt-oss-120b at 64 GB), but the guard sat on the automatic split, and upstream FlashML-org#418 replaced that with an explicit 'auto'. The split answers "more banks than we may pin"; --moe-bank-ram answers it first and differently, because the banks are a file mapping and only the resident prefix is locked. An explicit layer list or count still means what it says. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
MT-z
added a commit
to MT-z/FreeToken
that referenced
this pull request
Sep 16, 2026
…org#418), without FlashML-org#337 FlashML-org#418 replaced the per-family quantization branches with QuantConfig / QuantScheme / QuantMethod, and moved expert-bank ownership into the MoE kernel. 16 files conflicted; they were not 16 equal problems. Five were mechanical: * engine/cache_budget.py -- take FlashML-org#418's max_slots parameter (it replaces the hardcoded 992 marlin cap) and keep our +1 dummy-page reservation on the KV floor, with its test. * engine/config.py -- keep the fields, take the renamed-flag wording. * kernel/triton/fp8_pertensor_linear.py -- FlashML-org#418 removed the BaseOP layer, so only the round_e4m3 import is still needed; the e4m3 rounding fix itself is untouched. * models/qwen3_5_moe/model.py -- FlashML-org#418's ParallelLMHead takes a quant_config and selects the NVFP4 W4A16 kernel itself, which subsumes our Nvfp4LMHead special case (and the ~1 GB bf16 dequant it existed to avoid). Take theirs, keep the vision tower line. The other eleven were all FlashML-org#337's NVMe disk tier, whose plumbing point no longer exists: FlashML-org#418 deleted the per-model load_nvfp4_expert_sources(..., disk_tier=) in favour of a one-line nvfp4_expert_spec() plus expert_pieces, so the tier would have to be re-attached to the kernel-owned layout rather than merged. MT chose to drop it and rebuild it fresh, so it is removed whole -- moe/disk_tier.py and its test, the three CLI flags, the three EngineConfig fields, and every call site in layers/moe.py, moe/offload_cache.py, moe/expert_banks.py and engine/engine.py. It is recoverable from 9834eb3 / e4b565c. Three files needed hunk-level care because other work shares them, and a wholesale --theirs would have dropped it silently: * moe/expert_banks.py -- take main, then re-apply FlashML-org#334's effective_memory_available() (cgroup-aware host RAM) on top. * models/qwen3_5_moe/weight.py -- the vision rename hooks auto-merged outside the conflict markers; resolve per hunk so they survive. * engine/engine.py -- take main for the two bank-binding hunks, keep ours for the --moe-collect-stats routing histogram (main still has OffloadMoeCache.collect_decode_freq). layers/moe.py and moe/offload_cache.py turned out to differ from main only by the disk tier, so both are now identical to it. The dead make_offload_moe_cache hook went with the disk-tier branch that read it; nothing in the tree implements it. MT's serving command still parses: --moe-backend offload folds into moe_strategy='offload' and every other flag is unchanged, with one new deprecation warning at start-up. Assisted-by: Claude Opus 5 Tests: the merged tree ran 10 failed / 1904 passed / 92 skipped; the pre-merge try/all ran 8 failed / 1905 passed / 60 skipped, and the two failure sets differ by exactly the two above, both since fixed and verified in isolation (4 passed, and the admission test does fail when effective_memory_available's result is discarded). The confirmation run of the whole suite after the fix was interrupted at 50%, so "8 failed after the fix" is derived from those two facts, not measured end to end.
MT-z
added a commit
to MT-z/FreeToken
that referenced
this pull request
Sep 16, 2026
…off and FlashML-org#424's FTW script fb7f732 (FlashML-org#427) refactor(quant): hand the checkpoint QuantConfig to the weight readers 505477a (FlashML-org#424) chore(scripts): ftw_hotfix.py FlashML-org#427 is what makes this worth merging now rather than later. It moves checkpoint_quant_config from engine/config.py to models/register.py and installs the resolved QuantConfig in a module-level global from EngineConfig.model_config. Nothing in upstream main reads that global yet, so today it changes no behaviour -- but it is groundwork on the same path FlashML-org#418 and FlashML-org#426 already moved, and this box's model (Ornith, qwen3_5_moe, MIXED_PRECISION) loads through exactly that code. Divergence here costs more at the next commit than the merge costs now. FlashML-org#424 comes along because it sits between us and FlashML-org#427 on main. It is a standalone repair script plus docs, no engine change, and there are no FTW checkpoints on this box, so it does nothing here either way. Verified on this box (results/20260910-merge-427.txt in freetoken-systest): same 8 failures by node ID, 1,910 passed; greedy output bit-identical over 399 tokens at temperature 0.0; same startup warnings, same expert path. Assisted-by: Claude Opus 5
MT-z
added a commit
to MT-z/FreeToken
that referenced
this pull request
Sep 16, 2026
…ne piece of it try/all had kept try/all never carried the disk tier: 5e4edfa resolved the FlashML-org#418 merge "without FlashML-org#337" by hand. So daily's revert (6600db0) is daily catching up to try/all, not the other way, and all twelve conflicts are the two branches having deleted the same code by different routes -- daily from the pre-refactor base, try/all through the quantization refactor. Resolved to try/all's side on all twelve; verified afterwards that nothing disk-tier survives anywhere under python/ or tests/. The merge is not a no-op, though. offload_kernels.py had kept FlashML-org#337's `materialize_count` parameter on _materialize_layer_gpu and _materialize_layer_kernel -- dead since the "without FlashML-org#337" resolution, because the only caller passes nothing and `count` therefore always equalled cache.num_experts. Taking daily's side there drops it. Semantically unchanged, three lines less. Assisted-by: Claude Opus 5
This was referenced Sep 16, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Checkpoint quantization goes through three layers instead of per-family format branches:
QuantConfig: parsesquantization_config(ModelOpt, compressed-tensors, block-fp8, mxfp4,hf_quant_config.json) and answersscheme_for(name)per projection.QuantScheme: the storage contract (weight / scale descriptors, roles); the kind (fp8_tensor,fp8_block,mxfp8,nvfp4,mxfp4) follows from it.QuantMethod+ kernels: everyLinear*,ParallelLMHeadand MoE layer takesquant_config, prefix, declares its tensors and selects its kernel at construction. MoE kernels own the expert bank layout,packandapply; the offload cache, the FTW converter and the CPU executor follow that layout.Removed:
quant_linear.py×2,glm4_moe/nvfp4_linear.py, theBaseMoeBackendfamily,moe/nvfp4_backends.py, the per-format expert bank providers, reader-side dequantization and the format / debug environment switches (GLM runtime fp8 included; FTW files converted under those switches need reconversion).Flags:
--quant-backend layer[.kind]=nameselects kernels (--nvfp4-backendforwards to it);--moe-strategyreplaces--moe-backend(old spelling stays as an alias);--moe-cpu-layers autofor Windows/WSL pin budgets;--moe-cache-autorespects the kernel's slot limit; NVFP4 MoE defaults to Triton, Marlin / b12x on request.Verified
AIME25 pass@3 (
tests/e2e/test_aime.py) on RTX PRO 6000, timings within noise of main:tests/models/test_quant_config.pychecks the method / kernel of every probed projection per local checkpoint andscheme_foragainst the stored tensors;tests/kernels/test_fp8_blockscale_moe.pychecks the block-fp8 MoE kernels per activation.