Repository navigation
Integrate finegrained kernels: SOTA performance and per-module quantization - #49658
Draft
IlyasMoutawwakil wants to merge 101 commits into
Draft
IlyasMoutawwakil wants to merge 101 commits into
IlyasMoutawwakil wants to merge 101 commits into
Conversation
The kernels now store gate|up row-interleaved (gate j at row 2j, up j at 2j+1)
rather than as two N-apart halves, so the integration follows:
- GPT-OSS already ships this order on disk, so its de-interleave converters go
away entirely: _deinterleave_gate_up_rows, its use in FineGrainedMxfp4Deserialize,
and the whole FineGrainedGateUpBiasDeinterleave op plus its registration.
- Checkpoints that ship gate_up stacked are interleaved once post-load by
interleave_gate_up_after_loading, ordered BEFORE swizzle_scales_after_loading
(the swizzle packs whatever row order it finds) and skipped for mxfp4. This
cannot live in conversion_mapping's MergeModulelist+Concatenate chain — that is
shared with non-finegrained MoE models.
- _apply_gate de-interleaves with a stride-2 split, mirroring the fused epilogue's
split_gate_up; the two must agree or fused and unfused stop being comparable.
- swizzle_scales_after_loading drops gate= and gates on the doubled extent
(n_rows % 128, i.e. N % 64), so GPT-OSS N=2880 now pre-swizzles instead of
falling back to an affine scale read.
Adds tests/kernels/test_finegrained.py (24 tests), covering the post-load interleave,
that mxfp4 is left untouched, dispatch/swizzle gating, and the MergeModulelist path.
tests/kernels/test_finegrained.py: 24 passed.
The kernels store gate|up row-interleaved, but ``transform_weights_for_mega_moe`` does its OWN gate/up interleave and so takes the stacked form. Interleaving those modules at load and undoing it at the boundary is a pointless round trip — and interleaving twice is not the identity, it is a different permutation, i.e. silent garbage rather than a crash. So modules bound for megamoe are simply never interleaved: the post-load pass skips them and records what the module actually holds in ``_gate_up_interleaved`` (the bytes cannot be inspected to tell, and ``set_experts_implementation`` can switch backends after load). ``setup_megamoe_weights`` reads that flag and raises if it is ever handed an interleaved module, rather than re-permuting on an assumption. The other DeepGEMM expert paths need no change: their grouped GEMM is layout-agnostic and the gate split happens in ``_apply_gate``, which already de-interleaves stride-2. Verified: triton backend interleaves and flags; megamoe backend is left untouched with no round trip. The raise itself is unverified here — DeepGEMM's JIT needs a CUDA toolkit >= 12.9 and CUDA_HOME is unset on this box, so no megamoe path executes.
# Conflicts: # src/transformers/integrations/hub_kernels.py # src/transformers/quantizers/auto.py
The memoized `(epilogue, gate_up_quantization, down_quantization)` triple bought nothing where it matters. Constructing the three dataclasses measured ~1.4us per layer (~85us per step at 61 layers), and only in EAGER decode — under cudagraphs the host path does not run on replay, which is the mode this integration deploys in. Against that it held hidden mutable state on the module, and cached a gate_up quantization carrying an `output_recipe` that silently mismatches any caller asking for the unfused form. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The kernels now take a per-expert bias as a `bias=` operand, so a biased model no longer falls back to the unfused two-GEMM form just for having a bias. `_kernel_epilogue` stops bailing on `has_bias`, and both host-side adds go away: `_apply_unfused_gate_up` keeps only the activation, `_finish_down` only the routing-weighted reduce. That also drops the `torch.sort` the grouped path paid to index the gate_up bias by expert-sorted row -- the kernel indexes by the tile's own expert id instead. An unsupported act_fn no longer costs the bias its fusion either: the bias is an operand beside the epilogue rather than part of it, so it fuses whether or not the GLU does, and `epilogue is None` goes back to meaning exactly one thing. The bias rides the same output-row axis as the weight and scale grid, so `interleave_gate_up_after_loading` already delivers it in the kernels' interleaved order -- no reordering here. Tests assert the operand reaches both GEMMs, plus a new case pinning that an unfusable activation still fuses its bias. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Conflict resolutions: - quantizers/auto.py: keep the finegrained claims on fp8/nvfp4/modelopt (pre-quantized + MoE experts; NVFP4HfQuantizer only quantizes on the fly), take main's new gguf entry. - integrations/__init__.py: union — the finegrained module plus main's FP8Embedding exports. - integrations/deepgemm.py: ours (to_local unwrap + the stacked-gate guard); re-add the to_local import main's shim conversion dropped. - integrations/tensor_parallel.py: main's shim, with to_local restored on it — the kernel integrations still import it from this path. - distributed/sharding_utils.py: re-apply the 0-dim scalar guard (ModelOpt per-projection weight_scale_2/input_scale) inside the rewritten DtensorShardOperation.shard_tensor: replicated on the dense path, expert-ownership-filtered on the MoE path. - tests/tensor_parallel: main's rewrite, with the 0-dim sharding test re-added against the new engine (passes). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e kernels' MoE chain
The experts forwards are adapters over the kernels' `moe_fused_batched/grouped`: `_moe_operands`
hands them the module's tensors and the activation — a `get_supported_act_fns()` name the gate_up
epilogue fuses, else the module's own GLU as a callable the kernels run on the host, so any
activation works without a kernel release. The kernel bundle is loaded by symbol name
(`matmul_2d/batched/grouped`, the two forwards, swizzle/unswizzle, `get_supported_act_fns`).
Every layout difference between a checkpoint and what the modules hold is a `ConversionOps` with a
reverse, attached by the quantizer, so `save_pretrained` restores the checkpoint bitwise:
- `FineGrainedInterleaveGateUp`: stacked [gate; up] rows -> the kernels' [g0, u0, ...] order,
skipped when the model declares `is_concatenated=False` (GPT-OSS; the flag travels through
`use_experts_implementation`, which stamps its defaults after `__init__`) or the backend packs
gate|up itself (DeepGEMM Mega MoE);
- `FineGrainedScaleContainer`: a scale into the dtype its module holds — the same bytes for a
uint8 UE8M0 container (MiniMax), an exact cast for float32 values (dsv4-flash-base) — and the
module records the shipped container so the reverse restores it;
- `FineGrainedSwizzleScales`: the SWIZZLE_32_4_4 artifact for modules that hold 5-D scales
(`FineGrainedExperts.__init__` allocates that shape for every non-DeepGEMM backend on SM100
with a group-scaled format and activation quant; whole 128-row/4-col blocks only);
- `FineGrainedPackedBlocks` regroups GPT-OSS `{proj}_blocks` into the packed weight, so the
blocks/scales format is two one-to-one converters and round-trips too.
Catch-all converters cover keys that already arrive under the fused names and dense scale keys.
The post-load hooks, `_save_to_state_dict` override and post-load dtype cast are gone.
Also: the eager per-expert loop reads a swizzled stack's `[e:e+1]` slice and passes bias, NVFP4
global and activation format (it dropped all three); expert biases take the model dtype; the
`mxfp4` format pins UE8M0 scales; `FineGrainedEmbedding` (FP8 table, per-tensor scale) for
`modules_to_convert`; DeepGEMM backends refuse modules holding swizzled scales or interleaved rows
before loading the kernel; `disable_deepgemm_on_multi_device` is public.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…heckpoint precision for quantized params `FineGrainedQuantize` reads the module's weight format: block-FP8 in torch (per-block E4M3, fp32 or power-of-two UE8M0 inverse scales, from the module's block and held scale dtype), MXFP8 / MXFP4 / NVFP4 through the kernels' row-wise quantizers in one launch per tensor (NVFP4 normalized by the per-tensor / per-expert global `amax / (6 * 448)`), and emits the scale (and global) in the layout the module holds — the container dtype, the swizzled artifact — since the loader runs the quantization op after the converter ops. Core loader: a parameter the quantizer's op is about to quantize is materialized in the checkpoint's precision instead of the empty parameter's storage dtype (int8 storage zeroed FP4 weights; float8 double-rounded the existing FP8 path). Also: `auto.py` groups the finegrained keys under one comment and drops unused NVFP4 imports; the quantizer's config type hint names `FineGrainedConfig`; DeepGEMM guards (`_assert_affine_scales`, `_assert_stacked_gate_up`, group-32 divisibility) run before the kernel load; expert biases take the model dtype; comment trims for the noisy-comments check; tests for every path. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ion/Epilogue bundle fields The kernels' dispatchers now take input_recipe/output_recipe strings, so the bundle no longer needs the Quantization and Epilogue config classes; the linears hand the module-level activation_format straight through as input_recipe and the _kernel_quantization adapter is gone. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…uantization/Epilogue bundle fields" This reverts commit 07dae50. The kernels keep the Quantization/Epilogue op-boundary configs as the dispatcher API, so the bundle carries them and the linears build kernel.Quantization again. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…kip the kernels arch check The kernels now speak the module's vocabulary: activation_format (None = the weights' format, "bf16" = weight-only) goes straight to matmul_2d / matmul_grouped / the MoE forwards, so the Quantization/Epilogue bundle fields and both adapter helpers are gone. hub_kernels: kernels >= 0.16 checks a build's declared archs against the device; the deep-gemm build declares only 9.0a although it JIT-compiles per device, so the mapping opts out via a new check_arch passthrough (only forwarded when False, so older kernels releases without the keyword keep working). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…(pr-1018) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…s per-expert output norm modelopt NVFP4 checkpoints ship a second-level global per projection and a calibrated `input_scale`; both now reach the kernels, in either of the two layouts such a checkpoint uses — one tensor per expert per projection (GLM-5.2) or one stacked tensor per layer (the fused vLLM layout). A gate|up stack arrives with two weight globals per expert, calibrated separately. One conversion op owns every global of a layer and merges them to the one-per-expert the kernels take: the stack keeps the gate's, and the up half's folds onto the down projection, whose weight global scales the expert output back up and whose calibrated input global moves the other way, keeping the requantized intermediate on the range the checkpoint calibrated. That is exact arithmetic on fp32 scalars, unlike rescaling the up half's e4m3 block scales. Many-to-many converters are now opt-in per op (`ConversionOps.supports_many_to_many`) rather than a private whitelist. A model states its per-expert output norm the way it states its activation: `use_experts_implementation(post_expert_norm="<name>")` plus an `_apply_post_norm` the class must define. A name the kernels implement rides as that name and is folded into their reduce; anything else is the module, called on the routed rows. The DeepGEMM experts paths, which have no slot for such a norm, refuse instead of dropping it. Weight-only modules hold no activation global, since nothing there quantizes activations: the parameters are not allocated, the converters for them are not emitted, and a calibrated checkpoint's `input_scale` keys are ignored for that run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
One guard decorator on the DeepGEMM experts forwards instead of four statements per body, and that file's comments cut to what is not already in the code. `prefers_deepgemm_linear` moves to `deepgemm.py`, where its measured verdicts belong, so the dispatcher just asks. `_FineGrainedModule.local(name)` is the one accessor every forward reads operands through, so the DTensor unwrap appears once rather than at seven call sites. `_WeightFormat` is public (it is in `resolve_weight_format`'s signature), the `TRANSFORMERS_FINEGRAINED_NO_SWIZZLE` debug hatch is gone (tests patch the predicate), and `_set_optional_parameter` replaces six register-or-assign blocks. `check_arch` is passed straight through: `kernels` is pinned >= 0.16, where it landed, so the compatibility shim was dead. The scalar-shard branch becomes `_owns_expert`, stated as the general fact it is rather than as an NVFP4 anecdote. `FineGrainedConfig`'s docstring opens with a format table. The frozen fp8 modules' `stacklevel=2` is dropped: at module scope it named importlib rather than the module being deprecated. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…they served Every checkpoint `quant_method` the two answer to — `fp8`, `mxfp8`, `mxfp4`, `nvfp4`, `modelopt` — already routes to `FineGrainedHfQuantizer`, so both carry the same DeprecationWarning the fp8 pair does. Auditing that supersession found one real gap. `dequantize=True` on a GPT-OSS MXFP4 checkpoint produced no converter at all: the `_blocks`/`_scales` keys were wired only for the quantized path, so they went unmatched, where `Mxfp4Config(dequantize=True)` had handled them. The finegrained chain now dequantizes them, bit-equal to `mxfp4.convert_moe_packed_tensors` — the implementation it replaces — which the new test compares against directly. Two shared ops needed repair to get there: `FineGrainedPackedBlocks` regrouped a sibling `_scales` entry as if it were blocks, and `FineGrainedDequantize` left the scale it had consumed in the chain. The same path on NVFP4 or modelopt would have folded one block scale and dropped the second-level global, handing back a weight scaled by `1 / global`. Neither frozen integration supported that either, but it failed quietly — including where `dequantize` turns itself on (no GPU, compute capability < 8.9). It raises now. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # src/transformers/integrations/deepgemm.py # src/transformers/integrations/moe.py
It had become `FineGrainedFP8Config = FineGrainedConfig`, so the name still imported but the class was gone. That makes "frozen for backward compatibility" a fiction — the old name silently gained the new behaviour (`activation_format`, the modelopt payload parsing, the wider `post_init`) — and it rendered the same class under two `[[autodoc]]` headings. Restored verbatim, with a docstring line pointing at the config that serves the whole family. Splitting them means the two `isinstance` sites in `auto.py` no longer cover the new config by accident. `LOADING_ATTRIBUTES_CONFIG_TYPES` is the load-bearing one: it is what carries `dequantize` / `modules_to_not_convert` from a config passed to `from_pretrained` onto the checkpoint's own, so leaving it alone would have quietly stopped `FineGrainedConfig(dequantize=True)` from taking effect on every fine-grained checkpoint. Both configs are in it now. `quant_method="fp8"` still builds `FineGrainedConfig`: the frozen class is what you get by naming it, not what a checkpoint resolves to. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`c66564416e` restored the class and added a line saying it "configures the frozen `finegrained_fp8` integration". That is not true: it sets `quant_method="fp8"`, which `AUTO_QUANTIZER_MAPPING` routes to `FineGrainedHfQuantizer` — the frozen quantizer is not reachable through `from_pretrained` at all. Removing it also leaves the class byte-identical to main, which is what a class frozen for backward compatibility should look like in a diff. Which config serves which family is already stated where a reader meets it, in `FineGrainedConfig`'s own docstring on the same docs page. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three CI failures on a CPU-only runner, three causes: `post_expert_norm_name` is set by `use_experts_implementation`, but the shared forwards are the ones every backend adapts onto — including duck-typed stand-ins built without the decorator. The norm is optional, so read it as one. `_assert_stacked_gate_up` already read its second attribute defensively; `has_gate` now matches. `prefers_deepgemm_linear` calls `is_sm100()`, and `get_device_capability()` asserts rather than answering False on a CPU-only build. Caught there rather than gated on `is_available()`, so that a test faking the capability alone still answers — which is the convention the existing deepgemm tests are written to. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A converter whose target no module holds is a load failure, and the activation global is exactly that under a weight-only run: the experts allocate one only when they quantize activations, so the declarations are the only thing standing between the two runs. Checked against a real module's parameter names in both. Verified it catches the failure by removing the conditional: the weight-only run then converts `gate_up_proj_input_global_scale` onto a module without that slot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The integration reads better as `is_available() and capability >= 10`; the reason it wasn't was that three test sites faked the capability alone, which that form ignores. That is the tests constraining the integration, so the tests move instead: they now patch availability alongside the capability, and the module docstring says so. Replaces the try/except that was catching what the availability check now prevents. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… name A sharded parameter is already local by the time a forward runs, so the unwrap was doing nothing: FSDP2 unshards in its pre-forward hook, and a TP/EP plan swaps each DTensor for its local shard around the call (`MoeExpertsParallel` -> `_use_local_dtensor_params`). Verified both: under `fully_shard` alone a param reads as a plain Parameter inside forward, not a DTensor. So the invariant belongs to the distributed layer that wrapped the parameter, not to every backend that reads one. That removes `to_local` and the `local()` accessor. Absent slots still read back as None because `_set_optional_parameter` registers them as None parameters, which is what the kernels take for a missing bias or global. `_moe_operands` also stops building its dict through a `both(suffix)` helper: the kernels' arguments are now written out, so `down_proj_scale_inv=` can be found by grep. The five `getattr`s that remain are load-bearing — the kernels always say `gate_up_proj`, the module says `up_proj` when the model has no gate. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
IlyasMoutawwakil
changed the title
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
This was referenced Oct 16, 2026
Open
Closed
3 of 4 tasks
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: deepseek_v4, glm_moe_dsa, gpt_oss, llama4, minimax_m3_vl, ministral3, mistral4, muse_spark, qwen4_exp, blockfp8, finegrained, finegrained_fp8, mxfp4, mxfp8, nvfp4 |
Contributor
CI recapDashboard: View test results in Grafana |
This branch has not been deployed
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.
Integrates finegrained-kernels for state-of-the-art quantized dense and MoE inference, and redesigns quantization around fine-grained, per-module control. The kernel bundle covers full precision, FP8, MXFP8, MXFP4, and NVFP4 with fused MoE execution, activation quantization, and hardware-specific dispatch. Benchmarks and reproduction instructions compare real model shapes on B200 against existing kernel backends.
FineGrainedConfigmaps module patterns to format configs, so experts, dense layers, and other supported modules can use different weight and activation formats in one model. This includes DeepSeek-V4’s MXFP4 experts alongside block-FP8 linears.MIXED_PRECISIONassignments. Map supported FP8, MXFP8, MXFP4, and NVFP4 variants; reject unsupported runtime transforms and scale encodings rather than silently treating them as NVFP4.FineGrainedFP8Configintegration, which remains available explicitly.Validation: 278 CPU tests passed; 248 GPU integration/kernel tests passed with two skips; all 16 mixed-format load/save/reload cases passed against the updated local weight-only kernel bundle; 10 GPU generation/distributed tests passed with one expected DeepSeek-V4 compile skip. Static-cache decode compilation runs under distributed EP/TP. The older large-checkpoint Llama/GPT-OSS test classes were excluded from the GPU integration run.
make typingandmake stylepass.Kernel follow-up: explicit autotuning phase controls, native-UE8M0-optional imports and weight-only block-FP8 require updated finegrained kernel bundles. The embedding API is required; kernels release before or alongside this PR. CSF-compressed scales remain unsupported. NVFP4 PLE embeddings now load packed and decode only selected rows, with byte-exact sharded save/reload; full-table embedding dequantization is explicitly unsupported. Fused FP8/NVFP4 embedding lookup is integrated and requires the accompanying kernel release; compiled FP8 keeps its faster PyTorch expression. The Muse-Spark tiny fixture was rebuilt in the ModelOpt layout with the original NVFP4 exclusion policy.