[staging CI] unslothai/unsloth#10731 - #1411
Open
danielhanchen wants to merge 109 commits into
Open
danielhanchen wants to merge 109 commits into
danielhanchen wants to merge 109 commits into
Conversation
Refactor only, no behaviour change: the shipped preference table is empty
and every existing ladder test passes unchanged.
Two extractions, both because the next commits need to read the same
numbers the runtime uses:
- divisible_for_scheme(scheme) replaces the inline {fp8: 16, nvfp4: 16,
mxfp8: 32} map in quantize_transformer. The offline builder and the
checkpoint validator have to agree with the runtime filter here: a
checkpoint built at a different alignment floor holds a different set of
quantized Linears, and nothing downstream would notice.
- _auto_scheme_order(family, device, cap) is the walk that
select_transformer_quant_scheme and auto_scheme_candidates each used to
do for themselves. They now share it, so the winner and the candidate
list come from one order by construction rather than by two copies of
the same loop staying in step.
On top of that walk, _AutoPrefer / _FAMILY_AUTO_PREFER: a per-family head
tried ahead of the global _AUTO_LADDER tier, dropped when the capability
is below the row's floor or when the GPU is consumer-class and the row
does not say consumer_ok (the measurements behind a row are taken on
datacenter Blackwell). The deny list and the smoke probe still filter the
result, so a head can reorder what auto tries and can never widen it.
The table ships empty. Rows are added per family only once this PR's
promotion gates pass for that family.
torchao's NVFP4 dynamic-activation path takes its global activation scale
from torch.max(torch.abs(x)) over the WHOLE input, and max() over an empty
tensor raises. Micro-repro on a B200 (torch 2.12.1+cu130, torchao
0.17.0+cu130), a quantized Linear(1152, 3072) built with the runtime
config quantize_ uses, saved as scripts/nvfp4_zero_row_repro.py:
weight type after quantize_: NVFP4Tensor
M=1: ok shape=(1, 1, 3072) finite=True
M=7: ok shape=(1, 7, 3072) finite=True
M=5: ok shape=(1, 5, 3072) finite=True
M=0: RuntimeError: max(): Expected reduction dim to be specified for
input.numel() == 0. Specify the reduction dim with the 'dim'
argument.
So small M is fine (to_blocked pads the scale rows; there is no _int_mm
style row floor here) and only an EMPTY batch fails.
HunyuanVideo-1.5 reaches that state on every default t2v render. The
attention trim in diffusion_attention.py replaces the all-zero image
stream with image_embeds[:, :0], and the image embedder's two admitted
linears (1152 -> 1152 and 1152 -> 2048) receive it; the same trim can take
the optional byt5 stream to zero tokens, which reaches context_embedder_2's
three (1472 -> 2048, 2048 -> 2048, 2048 -> 2048). All five clear the nvfp4
filter (>= 512 features, both dims 16-aligned), so all five are quantized.
Fix, in the shape of the existing small-M padding:
- ZeroRowSafeLinear in diffusion_quant_pad.py answers an empty activation
itself, at the projected width plus bias. Bit-equal to F.linear, which
is empty at that shape too. Attribute and state-dict transparent, so the
checkpoint keys are unchanged.
- zero_row_tokens_for_scheme(scheme, family) next to pad_tokens_for_scheme.
nvfp4 only: fp8's per-row amax reduces along the feature axis and is
well defined for zero rows, mxfp8 blocks along the same axis, and int8
has its own answer through PadToMinM.
- apply_zero_row_guard called right after apply_small_m_padding on both
paths that put quantized weights in place, the runtime quantize_transformer
and the prequant load_prequantized_transformer.
Every admitted linear stays quantised. This is a shape guard at one point,
not a per-layer policy.
…ublishing on a second build scripts/build_prequant_checkpoint.py could only bake image DiTs: it resolves --family through diffusion_families.detect_family, which answers None for every video family, so the Wan and HunyuanVideo denoisers had no offline path at all. It also loaded a fixed "transformer" subfolder, which is one of the two experts a Wan2.2 A14B checkpoint has. The family lookup now asks the image registry first and the video registry second (--modality pins either), and --component names the denoiser subfolder to quantise. The component is stamped into the metadata and into the published filename (Wan2.2-T2V-A14B-transformer_2-NVFP4.pt), because both experts share a family, a scheme, a base and a key set, so the name is the only thing that keeps expert 1 out of expert 2's slot. Plain builds now publish under the derived <Model>-<SCHEME>.pt, which is the name the loader asks for first; the legacy transformer_<scheme>.pt stays resolvable as its fallback. Two correctness holes come with it: - the offline filter omitted require_divisible, so an fp8 or nvfp4 build baked the ragged linears the runtime leaves dense, and the mismatch surfaced only at the first real matmul of the first render. The builder now passes the runtime's own divisible_for_scheme and records it. - a corrupted build was indistinguishable from a good one. Every checkpoint now carries an md5 of each quantized weight's packed payload (qdata, scale and the per-tensor scale for nvfp4; the two int8 wrappers descended into), and --verify-against diffs this build against a second, independent one, printing each differing fqn with its parsed position and exiting 3. --upload-repo is refused without a verify that matched, and comparing an artifact with itself is refused too. No escape hatch: an unverified artifact is indistinguishable from a verified one once it is hosted. scripts/prequant_fingerprint.py prints or diffs the block for the artifacts that already exist. The format tag stays v1: the fingerprint is additive metadata that an older build ignores.
…filter floor before loading it Every check the loader made read what the artifact SAYS about itself. The bytes it holds were assigned into the served model without one of them being looked at, and a flipped byte in a quantized weight renders plausible garbage rather than raising: a truncated cache entry, a bad upload or a mutated hosted file all pass format, scheme, base, min_features and the exclusion set. The checkpoint now carries an md5 of each quantized weight's packed payload (written by scripts/build_prequant_checkpoint.py), and the loader recomputes it after validation and before the model is built. A mismatch logs the first five differing fqns and returns None, the dense fallback every other refusal here takes. UNSLOTH_PREQUANT_FINGERPRINT picks how much is checked: full (default, about 5 s on a 5B DiT), sample (a stable 1-in-8 by md5 of the fqn, so two loads check the same subset), off. It fails soft where softness is right -- an artifact with no block is accepted, since every hosted fp8 and int8 checkpoint predates it, and so is a build that recognises none of the weights, because a torchao payload rename should not refuse a whole scheme. Two metadata checks come with it, both accepting an absent field: - require_divisible. Same scheme and min_features, a different admitted set: an artifact baked before the builder passed the GEMM tiling floor holds the ragged linears the runtime leaves dense, and that surfaces at the first real matmul of the first render. - component. A MoE video family's two experts share family, scheme, base and key set, so every other check passes on the wrong one. load_prequantized_ transformer takes the component the caller is bringing up and refuses a checkpoint stamped with another. prequant_repo_filename also strips -nvfp4 and -mxfp8, so unsloth/Model-NVFP4 resolves Model-NVFP4.pt rather than Model-NVFP4-NVFP4.pt.
# Conflicts: # studio/backend/core/inference/diffusion_transformer_quant.py
The loader compares a checkpoint's base_model_id against the base it is loading for, so an artifact built from a local directory carried a path no load would ever match. --base-id stamps the repo id (and resolves the family) while --base stays the directory the weights are read from.
--gptq-dir hands every admitted linear its calibrated bf16 weight before quantize_, so the packed 4-bit operand in the artifact is the corrected one rather than round-to-nearest. The correction touches the 4-bit weight only, which is the rule the campaign measured. Do no harm, per layer. The calibration pass records the Frobenius WEIGHT error, which GPTQ raises on every layer by construction (it minimises the output error, not the weight error), so scoring on it admits nothing: 0 of 303 on TI2V-5B, 0 of 403 per A14B expert, 0 of 781 on HV15. --gptq-score-mode check therefore decides on the held-out OUTPUT error measured through torchao's own quantiser, and --gptq-score-mode meta keeps the weight-error rule available. A layer with no score, or with no corrected weight on disk, is left at RTN and counted, so an artifact can always say which of its weights are calibrated and which are not. After quantize_ the packed weights are dequantised and compared with the corrections they came from, and the worst deviation plus the fraction of moved elements is recorded: if the quantiser and the calibration pass disagreed about the grid, the artifact would hold a differently rounded weight that nothing measured. --gptq-dir with --convrot-groupsize is refused from the arguments alone, the correction having been solved against unrotated activations.
The video loader has only ever quantised its DiT the runtime way: from_pretrained
materialises the dense bf16 denoiser and quantize_transformer rewrites it in
place. That pulls the full bf16 shards on every fresh machine and peaks at the
dense size even though the steady state is half of it. The modular MiniMax-H3
workflow already avoids both by seeding a hosted checkpoint before its components
are built; this does the same for the conventional families, through pipe_kwargs,
in the shape diffusion_te_prequant uses for the text encoders.
Two things make the denoiser harder than the encoder. A dual-expert MoE
(Wan2.2-A14B) has two denoisers that share a config and a key set, so the second
artifact is addressed through the task slot of prequant_filenames: a task-specific
name gets no filename fallback, and here that guard is the whole point, since a
fallback would load expert 1 as expert 2 and pass every check on the way. And
seeding is all-or-nothing, the same rule the runtime block applies ("engaged on
only N/M experts"): one component that will not load means none is seeded and the
load runs dense exactly as it does today.
- video_denoiser_prequant.py (new): denoiser_components,
denoiser_prequant_sources, denoiser_prequant_pipe_kwargs. Pure orchestration,
torch-free at import, each expert released before the next is opened (an A14B
expert deserializes from a ~7 GB pickle).
- video.py: _video_auto_denoiser_scheme resolves the scheme before the download
plan and is pinned into the load, so plan and load can never disagree about a
pull that has already dropped the dense shards. The memory plan prices the
transformer term AND the build peak at the artifact's size, and re-plans plus
re-refuses at bf16 when nothing is seeded. _denoiser_prequant_covered and the
hub-file / cached-repo probes answer for every component; _base_download_files
takes skip_transformer_components, so both experts' weight shards are dropped
and both config.json stay. The runtime quantise is skipped for a seeded load and
the resolved record names the checkpoint.
- video_families.py: prequant_resident_gb_by_scheme and
video_family_prequant_resident_gb (per-scheme row, then the existing float).
No nvfp4 rows yet; those land with the artifacts.
- _VideoLoadState.h3_denoiser_pinned is now denoiser_pinned.
…e backend flashinfer 0.6.6 registers no torch custom op of its own, so Dynamo tries to trace its Python entry point and fails; the two calls are wrapped here with fake impls that reproduce its allocation exactly, the transposes inside the op, and the one-element PDL barrier between the quantiser and the GEMM. Every launch runs inside torch.cuda.device for the tensor's own device, which an AST test over the module enforces: flashinfer installs no device guard and an unguarded launch takes the card down rather than raising. select_nvfp4_backend decides torchao against flashinfer by import, capability and a memoised guarded preflight rather than by a version table.
Whole-model NVFP4 pays on the video denoisers and does not on the image ones.
What the campaign measured holding up there is a MIXED assignment: a small named
set of layers at 4 bits, the rest at fp8. This is that assignment, as data.
A policy names its layers by exact dotted suffix (never substring: flux's
single-block norm.linear must not select the double blocks' norm1.linear or the
top-level norm_out.linear) and carries a count for every rule, for the layers it
admits below the shared filter's floor, and for the final assignment. All of them
are asserted, and a mismatch raises PolicyMismatch rather than warning: every
count was measured on one base repo at one diffusers version, so a rule that now
selects a different number of layers is a different model, and building it anyway
produces an artifact whose quality nothing has measured.
Resolution is keyed on (family, BASE), through canonical_base so a mirror is the
same weights. An unnamed base resolves to None even where the family has exactly
one policy today: "the only one" is a fact about the table at this commit, not
about the model the caller is holding, and inheriting a verdict is the failure
this module exists to prevent.
quantize_with_policy runs two quantize_ passes over disjoint fqn sets, NVFP4
first. The order is what lets pass 2 also require a plain nn.Parameter: after
pass 1 the 4-bit layers no longer hold one, so no fqn-set bug can quantise a
layer twice. Both configs come from _make_quant_config, so they are the runtime's
own and go through _quiet_config. A post-pass walk then asserts each layer's
weight class by name (NVFP4Tensor / Float8Tensor / Parameter), because a torchao
that silently declines a layer would otherwise leave an artifact whose metadata
claims a precision it does not have.
Tables, with the counts verified this way: each family's diffusers transformer
class instantiated on the meta device from the real config (Z-Image-Turbo and
Qwen-Image from the local mirrors, FLUX.1-schnell from its hub transformer
config), walked for nn.Linear, filtered with make_filter_fn(512, ("lora_",),
require_bf16, require_divisible=16), then assign_precisions run over it. Admitted
sets 239 / 499 / 843 reproduce the linear census the campaign measured on.
- zimg_f8mod_toq34_v1: 34 attention.to_q at nvfp4, 237 fp8 (including the 32
(256, 15360) adaLN_modulation.0 the 512 floor rejects and this policy admits),
5 bf16.
- flux_mod_single_v1 (schnell only): 38 single-block norm.linear at nvfp4,
461 fp8, 3 bf16. The plan's draft said 0 bf16; x_embedder and proj_out are
64-wide and the timestep embedder's first layer is 256-wide, so it is 3.
- qwen_p02_v1: 60 img_mod.1 + 60 txt_mod.1 at nvfp4, 723 fp8, 3 bf16.
No family may lower min_features to reach its modulation projections: 256 would
also admit t_embedder.mlp.*, which cannot be quantised at all because
TimestepEmbedder.forward reads mlp[0].weight.dtype. That is why the floor is
overridden per named layer with its exact shape asserted instead.
Tests build census-exact synthetic trees (every fqn and every (in, out) pair off
the census, as meta-device nn.Linear so a 60-block DiT costs no memory) and
assert the totals, the resolution matrix, a rename, a shape drift, a width change
that only moves the totals, suffix-not-substring, t_embedder.mlp bf16 under every
policy, the two passes' disjointness and order, and the post-pass walk.
A per-layer policy artifact is a mixture -- some layers NVFP4, the rest fp8, chosen by a table solved on one base repo -- and nothing about its weights says which layers got which. A build predating diffusion_nvfp4_policy reads such a file as a whole-model nvfp4 checkpoint: it loads clean, it renders, and the precisions are not the ones any gate measured. So it gets its own on-disk tag, v3, biconditional with the declaration exactly as v2 is with the rotation. - PREQUANT_FORMAT_POLICY joins PREQUANT_FORMATS, and prequant_format_for stamps it. A build declaring both a rotation and a policy raises instead: there is one tag slot and two things an older build has to be warned about, so whichever tag it got would say the other half is absent. - _validate_policy mirrors _validate_activation_rotation and refuses five ways the artifact and this build can disagree: the tag/declaration biconditional, a block that does not parse, a scheme that is not nvfp4, a family and base this build resolves no policy for, and a declared (policy_id, policy_version) or count set that differs from the in-tree table's. A retuned table bumps its version precisely so the artifacts built under the old one stop loading rather than being read as the new one. - _SCHEME_REQUIRED_GLOBALS["nvfp4"] becomes the union with fp8's names. A policy checkpoint's state dict holds Float8Tensor weights beside the NVFP4Tensor ones, so answering on the 4-bit names alone would report the file loadable and then fail mid-unpickle, after the plan had dropped the dense shards. Every torchao shipping the prototype nvfp4 tensor ships the fp8 one. - _fp8_activation_floor_present now skips tensors that are not Float8Tensor (by class name, so the check stays torch-free). It returned on the FIRST tensor carrying act_quant_kwargs, and an NVFP4Tensor carries one too, with no hp_value_lb because its activation quantiser has no such knob -- so a mixed dict whose first quantised weight is 4-bit was refused for a floor that layer neither has nor needs. That is every policy checkpoint of z-image and qwen. The two stand-in tensor classes in the prequant tests are now one shared pair (Float8Tensor, NVFP4Tensor) since the floor check reads the class name; the convrot format-set assertion gains v3.
…chao payload
torchao's qdata and swizzled scale ARE flashinfer's packed weight and 128x4 block-scale buffer, so
a hosted NVFP4 artifact moves onto the fast GEMM as a rename rather than a requantisation: the
block scales are bit-identical and the 4-bit codes agree on 99.8 percent of nibbles, every
disagreement being one e2m1 step where the two encoders break a tie differently. The layer takes a
BAKED activation global scale and refuses conversion without one, which drops the canon layer's
running-min calibration state machine along with the black-frame latch it caused. Measured on a
B200: the converted layer sits the same 0.131 from the dense Linear as torchao's own module at
every token count in {1, 512, 4096, 16384}, M=1 runs 0.52 ms on torchao against 0.16 ms here,
a two layer block compiles fullgraph in one graph with no breaks, and a captured graph replays it.
--policy auto (the default) applies the policy diffusion_nvfp4_policy resolves for (--family, --base-id or --base) when --scheme is nvfp4, and builds the whole-model artifact when nothing resolves. Every existing invocation therefore keeps building exactly what it builds today: no policy describes an fp8 or int8 artifact, and none resolves for the video families. --policy off forces the whole-model build, and --policy <policy_id> is a PIN rather than a lookup: it must be what resolves for the same family and base, so an operator asking for the layer set they measured gets a refusal when the table has moved under them instead of a differently-quantised artifact. resolve_build_policy is a pure function so all of that is asserted without a GPU. With a policy the single quantize_ call becomes assign_precisions (which is also what scopes the GPTQ corrections and what raises on a diffusers rename) plus quantize_with_policy's two passes, and the metadata gains the nvfp4_policy block, so prequant_format_for stamps v3. The fp8 half of the build bakes an accumulate mode and a per-row granularity in like any other fp8 build, so both are recorded. --gptq-dir keeps working and is scoped to the layers the policy assigned to nvfp4, never the whole admitted set: the campaign measured the correction on the 4-bit operand ALONE, and a correction that also becomes the source of an fp8 replica raised the error 46 percent. A static policy gives that by construction, but only if the set it is applied to is the policy's. Refused before the load, from the arguments alone: a policy together with --convrot-groupsize (both rewrite the weights before quantize_, both claim the one format tag slot, and the rotation was solved for one quantiser over the whole model), a named policy that does not resolve, and a named policy on a scheme that is not nvfp4. policy_expected_counts drops zero entries before the assignment is compared: a Counter never records a precision no layer took, so a table spelling out "bf16: 0" would otherwise fail against an assignment that is exactly right.
FlashInfer's cutlass FP4 kernels take the stream from the tensor and install no CUDADeviceGuard of their own, so a launch made while the current device is not the tensors' device runs against a foreign context: the kernel hangs and the card ends up in "GPU requires reset". Three cards on this host were lost that way. This walks the list and writes it down. The audited sites are the quantise body, the GEMM body (one guard spanning the ordering barrier and the GEMM, because the barrier only orders anything if it fires on the card the GEMM reads from), the preflight, the prewarm entry and the traced layer forward. The layer forward keeps its guard under torch.compile rather than skipping it under is_compiling(). A live context manager in a traced region is a plausible graph break, so it was measured instead of assumed: a two-layer NVFP4 block still compiles fullgraph to one graph with zero breaks on torch 2.12. The source guard test grows two teeth. It now walks the AST instead of grepping for a substring, so an aliased cuda.set_stream is caught and a local named reset_stream_cache is not, and it bans set_device alongside set_stream: set_stream silently sets the current device, and set_device moves the value the guard restores. New tests/test_diffusion_nvfp4_speed.py proves the property rather than the syntax. A stubbed torch turns torch.cuda.device into a recorded push/pop and every stubbed flashinfer call logs the device it saw, so with tensors on device 1 and the current device 0 the quantise, GEMM, preflight and prewarm paths are each asserted to launch on device 1 and to leave the current device at 0, on a one-GPU host with no flashinfer at all. A guard leaked by a raising launch is caught too. The real two-card version is there and skips itself when device_count() < 2. 142 passed, 1 skipped (the multi-GPU case) on a single B200.
…re capture An explicit ``nvfp4`` on an image base with no hosted checkpoint used to quantise the whole model to 4 bits, which is not what the campaign gated: the gated model is a named set of layers at 4 bits over an fp8 one. ``quantize_transformer`` now takes the upstream ``base_repo``, resolves the policy for ``(family, base_repo)`` and, when one exists, takes the same ``assign_precisions`` + ``quantize_with_policy`` path the offline builder takes, recording the id on the transformer. With no policy (an unnamed base, a video family, any other scheme) nothing changes. Two capture-safety items go with it. ``GraphedForward`` autotunes the FlashInfer NVFP4 GEMMs in its warm-up, before the capture, since the autotuner profiles by launching candidate tactics and under capture those would be recorded rather than measured; the token counts come off the call's own input shapes, which is the only place the render's resolution is known. And it refuses to capture a module holding a layer whose activation scales are not baked: a scale still being calibrated would be frozen into the graph at whatever value the capture happened to see. The loader prewarms the M = 1 modulation shapes after a pre-quantized load, which is all it can know without a resolution.
…eone measured nvfp4 is the one scheme whose quality is a property of one artifact rather than of the scheme: a per-layer policy on the base it was solved against is a measured model, the same token on an unmeasured base is a guess. So the ladder does not get to ask "can this GPU run fp4"; it asks whether a reviewed record says the 28-pair accuracy gate passed for that family on that base at the policy this commit resolves. The record is data, checked in beside the policy tables, written only by scripts/record_nvfp4_gate.py from a gate results.json and only into a reviewed commit. The verdict is re-derived rather than trusted: nvfp4_gate_passed resolves the in-tree policy and compares (policy_id, policy_version) against the record's, so retuning a policy invalidates every verdict taken on the old one instead of carrying it onto layer precisions no gate ever saw. Ships with an empty records list, which means every answer is False and nothing changes yet. The writer fails closed on what it can check without a GPU: the policy must be the one the tree resolves for that base, the checkpoint's sha256 is computed here rather than taken on trust, a run that did not pass is refused unless it is explicitly recorded as a failure, and a key the file already carries is refused rather than duplicated.
…ied for training Wires the gate record into the three places that decide what nvfp4 means. The ladder. The commented Blackwell tier is gone; nvfp4 reaches auto only through the per-family preference table, whose rows may now declare themselves gated: the head applies only where a reviewed record covers THIS base at the policy this commit resolves. The three image rows land with it, ordered (fp8, nvfp4, mxfp8, int8) so nvfp4 can never displace fp8, and inert as shipped since no record ships. A scheme in require_prequant (nvfp4) is additionally offered under auto only when the caller's has_prequant probe says a usable checkpoint exists for this load: the gate measured the hosted artifact, not an on-the-fly build. An explicit request is unaffected. The deny. _family_denied takes the base and lifts the qwen nvfp4 entry when the gate passed for that base alone, since the deny was measured on whole-model nvfp4 and a per-layer policy checkpoint is a different model; the refusal message now names the missing record rather than claiming a measured breakage. Training denies nvfp4 for every family as a rule: the evidence is an inference gate on a frozen forward, and nobody has measured a LoRA over 4-bit frozen linears. The size. A policy base is sized by its policy factor rather than by whole-model 0.33 -- the artifact is mostly fp8 by weight -- in both estimate_dense_quant and the per-family info summary, so the planner and the UI stop promising gigabytes the load then takes. base_repo is threaded through the selector, the candidate list and the call sites that know it (the dense build, the auto prequant retry, the dense candidate resolver, the video seed probe).
… scales Two passes over the DENSE pipeline, before anything is quantised, on the layers the build takes to 4 bits (a policy's NVFP4 set, or every admitted linear for a whole-model artifact). --gptq-prompts N accumulates a bounded per-layer Hessian on sampled denoise steps, corrects each weight onto the NVFP4 grid and keeps the correction only where it lowers that layer's output error scored through torchao's own quantiser, which is the one that packs the artifact. The damping escalates and then raises: a layer whose Cholesky will not factor is a failed calibration, not a reason to ship one weight at round-to-nearest with nothing saying so. --gptq-max-regressions is the deliberate escape hatch for applying corrections against their own score, least harmful first. --bake-activation-scales measures each of those layers' activation amax over every step of the calibration prompts and stores a_gsf = 6 * 448 / amax per layer. The flashinfer NVFP4 backend refuses an artifact without them, because learning the scale at run time is neither capture-safe nor deterministic and is the mechanism behind the flux black-frame latch. The two flags are independent; the bake alone is a supported build. Both passes are seeded per prompt so a second build reproduces the same corrections byte for byte, which is what the publishing gate compares. The calibration prompts are a reviewed file, disjoint from the accuracy gate's evaluation suite, and what an artifact was calibrated on is recorded in its metadata.
… one per GEMM FlashInfer launches the cutlass FP4 GEMM with PDL while the CUTLASS griddepcontrol instructions that make PDL safe are compiled out of its build, so a kernel must exist between the activation quantiser and the GEMM or the GEMM reads operands the quantiser has not finished writing. What protects it is a kernel EXISTING, not that kernel writing M x N bytes: on the trigger harness a bare allocation fires 50/50 while a one-element kernel is 0/50, exactly like the full memset. The GEMM op was getting that kernel from a per-call torch.zeros(1), which allocates a buffer purely to obtain the fill that comes with it. This keeps one 1-element bf16 buffer per device and fires zero_() on it. The buffer is never read, by this module or by the kernel, so the forward-33 objection that argued for a fresh buffer every call does not reach it: that was about caching the GEMM's OUTPUT, where one transient NaN write latches into every later render at the same token count. Never allocated during a capture. An allocation made inside a capture comes from the graph's private pool and dies with the graph, so a cold key under capture pays an uncached buffer instead; the prewarm forwards leave every device warm before any capture, and the pointer is asserted stable across capture and replay. Not an op argument either: a mutable tensor input would go through auto_functionalized and be cloned per call, which is the cost this removes. UNSLOTH_NVFP4_ZERO_BUFFER=1 still restores the full memset. reset_nvfp4_state() drops the barriers and the tuned shape set on unload, next to the CUDA graph pool teardown in both diffusion.py and video.py, so a buffer from one model's allocator state is never inherited by the next model's capture. M1, B200, arms interleaved and alternated per round, host min and device median over 7 rounds x 50 iters, outputs/nvfp4_pr3/m1_barrier.json. The barrier timed alone, which is where a 1.7 us effect is visible at all: arm barrier host us barrier device us per-call zeros 8.09 - 8.20 8.32 - 8.53 persistent 6.40 - 6.52 6.68 - 7.10 full memset 7.49 - 7.67 10.29 - 107.60 so 1.65 to 1.71 us of launch saved per GEMM, and the memset arm costs +3.6 us at the small shapes and +100.6 us at M=16384 for no additional protection. Per GEMM end to end, per-call zeros -> persistent: zimage mlp-up M4096 74.23 -> 73.74 us, mlp-dn 72.68 -> 70.69, flux qkv 69.03 -> 68.73, mlp-up M16384 241.18 -> 237.41. Bit identity holds across all three arms at all four shapes, and over 50 iterations on real operands. 150 passed, 1 skipped on the nvfp4 + cuda graph + precision files; 467 passed on the video, unload and prequant files.
… epilogue mm_fp4 has no bias epilogue in flashinfer 0.6.6, so the bias is a separate pass over the M x N output, and profiling put that pass at 48 percent of layer time at M=16384/N=12288. It is not memory bound: a plain copy_ moves the identical bytes at 6.28 TB/s where the broadcast add_ manages 1.77, so it is the kernel, not the traffic. diffusion_nvfp4_bias.py flattens the add to a 1-D pass, accumulates in fp32 and rounds once on store, which is what torch's own bf16 add does. The fused CUTLASS bias epilogue is DROPPED rather than shipped. It is genuinely free in the GEMM (0.0643 ms against 0.0628 without the bias) but it is not bit-identical to the unfused path: max abs 2.0 to 4.0, because it adds the bias to the fp32 accumulator before the single rounding while the unfused path rounds to bf16 first. That breaks the zero-max-abs contract every precision and CUDA graph test here is written against, and it is sm_100a only behind a 104 s runtime nvcc build. Reopen on either of two conditions: flashinfer exposing bias= on mm_fp4, or the GDC fix landing so the epilogue can be validated against a rounding-equivalent reference. The CUTLASS sources stay in scripts/ as research artifacts. Eager only. Inductor already fuses the bias into the neighbouring GELU or residual (0.0156 ms at M=4096 by differencing bias=True against bias=False in a compiled layer), so an opaque launch would block the better thing; fused_bias_add_ returns add_ under torch.compiler.is_compiling() before it checks anything else. Eager is not a warm-up detail here: flux under the default dynamic=True fails quantized lowering and every scheme falls back to it. UNSLOTH_NVFP4_FAST_BIAS=auto|0|1, documented with the other NVFP4 knobs in diffusion_speed.py. The launch sits inside a device guard, because Triton takes its device and stream from the CURRENT context rather than from the tensors it is handed. The AST guard test now understands a Triton launch (a Call on a Subscript) and would fail on an unguarded one. M3, B200, min of 7 rounds x 20 iters, outputs/nvfp4_pr3/m3_bias.json: shape M x N add_ fused speedup zimage mlp-up M4096 4096x12288 0.1131 ms 0.0318 ms 3.56x zimage mlp-dn M4096 4096x3072 0.0252 ms 0.0158 ms 1.60x flux qkv M4096 4096x4096 0.0332 ms 0.0160 ms 2.08x zimage mlp-up M16384 16384x12288 0.4552 ms 0.1219 ms 3.74x torch.equal against add_ is True at all four, and at all seven verification shapes in the test (16384x12288, 4096x3072, 1024x12288, 333x4096, 1x512, 9304x3072, 7x1024). The measurement's third arm, a compiled add in ISOLATION, lands at 0.0351 to 0.1234 ms: it still pays its own launch, which is why the number that matters for the compiled path is the 0.0156 ms differenced one where the add has been fused away. 170 passed, 1 skipped.
…lowlist and a self-check
Block-compiled z-image 512px at 100 percent 4-bit spends 0.618 s of wall clock on
0.125 s of GPU work, 20 percent busy, against 51 percent for the same model in fp8
with the same 0.126 s of GPU work. Identical device work, 2.5x the wall clock: the
GEMM is not slow, the host cannot issue it. mm_fp4 rebuilds a four-entry dict of
runner-factory lambdas, a fresh cutlass runner and an AutoTuner cache probe on every
call, and fp4_quantize re-runs two device queries and an f-string module lookup, and
a render makes 2151 of those calls.
diffusion_nvfp4_dispatch.py caches (runner, tactic, workspace) per (M, K, N, backend,
device), the quantiser's bound pybind function per device, and the two weight
transposes, then calls straight through. Every private flashinfer import in the tree
is in this one file, in one try.
Three fences, because this is private surface. An EXACT version allowlist ("0.6.6",),
not a minimum, since a symbol that moves in 0.6.7 is not a bug in 0.6.7. One try
around the whole private import list, so a missing symbol takes the entire fast path
down rather than half of it. And a runtime verify(): quantise and GEMM both ways on
the device and require torch.equal, run once per device from the preflight, off the
request path, with nothing fast used anywhere until it passes. A symbol that still
exists but means something else is exactly what a version check cannot see.
UNSLOTH_NVFP4_FAST_DISPATCH=auto|0|1, where 1 skips the version check and nothing
else. Both op bodies branch fast or public inside the same device guard, and the plan
builder opens its own: _get_cache_buf allocates the workspace on the CURRENT device
and choose_one launches every candidate tactic while profiling.
Capture safe by refusing: gemm_plan returns None for a cold key while the stream is
capturing, so profiling launches can never be baked into a graph. The prewarm makes
every key warm first, so this is a fallback rather than a path. The transpose cache is
bounded at 4096 and keyed on (data_ptr, shape), so a reallocated buffer cannot return
a stale view, and reset() drops all of it on unload with the barriers: the views hold
weight buffers and would otherwise pin a freed model.
All the private symbols anom829_fastops.py used are present unmoved in the installed
flashinfer 0.6.6 (AutoTuner, get_fp4_quantization_module, DEFAULT_WORKSPACE_SIZE,
_MM_FP4_TUNING_CONFIG_128x4, _get_cache_buf, get_cutlass_fp4_gemm_module,
device_support_pdl, get_compute_capability); no adaptation was needed.
M2, B200, host min and device median over 7 interleaved rounds x 50 iters,
outputs/nvfp4_pr3/m2_dispatch.json. Host us per call, public -> cached:
shape mm host mm device quantize host
32 x 2560 x 3840 55.73 -> 19.14 60.34 -> 19.54 15.10 -> 14.09
1056 x 3840 x 3840 58.57 -> 19.49 62.00 -> 20.01 15.27 -> 14.37
1056 x 3840 x 10240 54.39 -> 19.03 59.26 -> 25.90 15.33 -> 14.26
4128 x 3840 x 10240 58.53 -> 19.60 67.96 -> 65.91 15.86 -> 14.75
wan5b 27280 x 3072 x 14336 59.54 -> 20.78 521.00 -> 519.38 15.89 -> 14.76
35.4 to 39.1 us of host time saved per GEMM, reproducing the 61.8 -> 18.1 us prior.
The device column is the control: where the call is launch bound it tracks the host
because the "device time" is the gap between launches, and where the GEMM is actually
compute bound (4128: 67.96 against 65.91; Wan 5B: 521.00 against 519.38) it does not
move. Same tactic (19 at the Wan shape, from flashinfer's own AutoTuner), same runner
object, torch.equal at all five shapes for both the quantiser and the GEMM.
202 passed, 1 skipped across the six nvfp4, cuda graph and precision files.
…running nvfp4 is the one scheme with two implementations, and which one served is decided per device at load time: an artifact without baked activation scales, a failed preflight, or a Windows host all leave the model on torchao under the same "nvfp4". So the scheme alone cannot carry a speed claim, and there was nowhere to read the difference from. status() now reports transformer_quant_backend beside transformer_quant, read from the module tree rather than from what the load intended (conversion is all or nothing, so one converted layer means it ran), null for every other scheme. It flows through the existing status route and the loaded-models row, which appends it after the precision only where there is one to append.
# Conflicts: # studio/backend/tests/test_diffusion_transformer_quant.py
Three reviewed entries in nvfp4_gate_record.json, from the 28-pair accuracy gate run on a B200 with CUDA graphs captured and the flashinfer backend (torch 2.12.1+cu130, torchao 0.17.0+cu130, diffusers 0.40.0.dev0, flashinfer 0.6.6). Every pair is scored against the dense bf16 reference of the same base, same seed, same schedule. z-image, zimg_f8mod_toq34_v1 v1, RTN build (sha256 d2f80a4a32c7): 24/24 pairs, LPIPS mean 0.156 max 0.529, SSIM mean 0.848, CLIP delta mean 0.0052. z-image, zimg_f8mod_toq34_v1 v1, GPTQ build (sha256 22771918e8ef): 24/24 pairs, LPIPS mean 0.141 max 0.577, SSIM mean 0.866, CLIP delta mean 0.0063. flux.1, flux_mod_single_v1 v1 (sha256 5b344cafa388): 24/24 pairs, LPIPS mean 0.181 max 0.427, SSIM mean 0.800, CLIP delta mean 0.0064. z-image runs 24 pairs rather than 28 because the distilled Turbo schedule has no CFG, so the high-guidance stress case does not apply; flux schnell is the same. qwen-image gets NO record. Its qwen_p02_v1 run passed every per-pair check but failed the suite aggregates (LPIPS mean 0.3094 > 0.25, CLIP delta mean 0.0141 > 0.01, 6 of 28 cases over the per-case CLIP cap), so the family deny stands and auto keeps refusing nvfp4 there. The same z-image checkpoint measured on torchao instead of flashinfer scores LPIPS mean 0.1555 / SSIM 0.8447 / CLIP 0.0049 against flashinfer's 0.1564 / 0.8475 / 0.0052, which is trajectory noise: the two backends are the same numbers, not the same pixels. record_nvfp4_gate.py gains --gate-script: the gate driver lives outside this repo, so the sibling default hashed nothing and the record could not name what produced it. test_every_shipped_record_names_a_checkpoint_the_family_hosts fails until the nvfp4 prequant_repos rows land with the upload commit, which is where the plan puts them.
… pipeline
--bake-activation-scales only ever ran an image pipeline: a square grid, a schedule looked up in
the image generation-defaults table, and one denoiser handed in as "transformer". A video family
answered all three wrongly and quietly. It has a third axis (the frame count changes what a layer's
activations reach as much as the frame size does), it carries its own schedule and guidance on its
registry row rather than in the image table, and a dual-expert family has TWO denoisers under one
name, so a build of the second expert would have handed its weights to the pipeline as the first
one and baked the wrong expert's amax into it.
* --calib-resolution reads as WxHxFRAMES for a video family (default 832x480x25) and stays a
square size or WxH for an image one, which is refused a frame count rather than silently
rendering a still. A frame count off the family's k * frame_step + frame_offset lattice is
refused before the dense load instead of deep inside the pipeline hours later.
* --calib-steps defaults to 20 for a video family. The passes measure a per-layer second moment
and a per-layer amax, both of which converge over the trajectory rather than over its size, and
the shipped 1280x704x121 at 50 steps is hours per artifact for the same numbers.
* The denoiser goes into the pipeline under its own component name, so --component transformer_2
calibrates the low-noise expert and the high-noise one loads from the base. The boundary
between them is the pipeline's own switch, so each expert is measured on exactly the steps it
runs and neither pass needs to know where the boundary is.
* HunyuanVideo-1.5 takes no guidance kwarg at all; its scale is set on the guider, as the runtime
path does it, rather than left at whatever the checkpoint shipped with.
* The baked-scale metadata records the grid it was measured at, and a prompt module declaring
CALIB (the investigation's own video calibration set, which the GPTQ weights were solved on) is
read alongside one declaring CALIBRATION_PROMPTS.
--gptq-dir and --bake-activation-scales compose unchanged: load dense, apply the corrections that
are measured to help, bake the scales on the corrected weights, quantize_, save.
…e top level convert_nvfp4_backend already serves both artifact shapes, because it keys on the per-layer scales being present rather than on a policy resolving; the CUDA test now proves that for a whole-model artifact (every admitted linear at nvfp4, no policy block), which is what PR 1's video checkpoints are and what the flashinfer backend has to serve for video promotion to be possible at all. The one place that still assumed a policy was the refusal message: _declares_baked_scales read the flag out of the policy block alone, so a whole-model artifact whose bake produced nothing was reported as one that never asked for a bake. Those are different problems -- the first is a build to rerun, the second is a backend to stop asking for -- and the log line is the only thing that says which. It now reads the top-level flag too.
…n scales --gptq-dir ran AFTER the calibration block, so a build combining it with --bake-activation-scales measured every layer's activation amax on the uncorrected weights and then replaced them. The scales an artifact ships would describe a model that was never built: the corrected layer's own output moves, and so does the input of every layer downstream of it. The two GPTQ sources are mutually exclusive, so moving the replay above the calibration block leaves the in-builder path (Hessians on the uncorrected weights, correct, then bake) exactly as it was -- that order is fixed for the same reason from the other side, and calibration_stage_order still asserts it.
--base-id exists so a local mirror can declare the repo it mirrors, but nothing enforced that: with a remote --base the weights came from one repo while the family, the published filename and the stamped base_model_id followed --base-id, so a 480p HunyuanVideo build could be published under the 720p identity and pass the loader's base check. A remote --base now has to be the same model as --base-id.
A video load superseded between its prewarm and its commit returned at the token check before the aborted-load reset, leaving its denoiser pinned by the transposed-weight views. The cache is keyed on the weight object through a weakref whose callback drops the entry, and holds a view of the detached tensor so the weakref can fire; the storage goes with the module, and no reset touches the replacement load's live entries. CUDA test frees the weight bytes after the layer is collected.
…nt preflight failures out of the cache, and report the video NVFP4 backend Only the flashinfer layer consults the step controller, so an armed schedule on a torchao load logged and counted protected steps that never ran; the controller now knows which live layers can take the branch and refuses to arm without one. The preflight memoised an allocation failure taken while the model about to be evicted still owned the card, dropping the NVFP4 rung for the rest of the process; allocation failures are answered but not cached. The video status now carries transformer_quant_backend like the image status, read from the module tree of every expert, and the frontend row shows it.
…ifact under the cancel event, and release the checkpoint before the device copy The download plan never asked whether an artifact-sized memory plan still offloads, but the load drops the seed under any offload policy, so the plan dropped dense shards the load then topped up inline; the plan now runs the same arithmetic, and the dense-encoder re-plan re-takes the seed decision. The seeded denoiser artifact was pulled inside load_prequantized_transformer with no cancel event; it is prefetched next to the encoder artifact. load_prequantized_transformer kept ckpt and state_dict bound across to(device), which on a unified-memory host doubles the transient peak; both are released first.
# Conflicts: # studio/backend/core/inference/diffusion_prequant.py # studio/backend/tests/test_diffusion_prequant.py # studio/backend/tests/test_video_backend.py
…celerator torch.device(int) consults the accelerator API and raises on a CPU-only host, so nvfp4_preflight builds the cuda device from an integer index itself; the prewarm test stubs flashinfer so the tuning loop it replaces is reached.
… only sampled fqns The plan runs while the previous pipeline is still resident and can decline the NVFP4 seed on memory; the load then re-took the decision against post-teardown capacity and fetched the artifact inline on top of the dense shards the pull had kept. The decline is now pinned into the load with a sentinel next to the pick. The builder records which base the GPTQ corrections were calibrated on and refuses a checkpoint whose base does not match. A covered payload class with no hashed attribute reads as uncovered instead of the digest of an empty stream, and sample mode fingerprints only the fqns it compares.
Record identity carries the checkpoint digest, so one policy can hold a failed row for one artifact and a passing row for another. The reader validated only the first match, so a bookkeeping failure recorded before the passing artifact hid the pass and the gated checkpoint loaded at the wrong precision.
One policy can hold an artifact gated on torchao next to one gated on flashinfer; the backend check read only the first passing row, so a device on the other backend lost nvfp4 from auto depending on file order.
danielhanchen
force-pushed
the
pr-10731-ci
branch
from
September 14, 2026 00:47
f7fc95b to
6abb382
Compare
Measured eager on B200 against Tensor.add_, bit-identical at every shape: the kernel is 0.26x to 0.37x below 4 M elements (its launch is 20 to 28 us against 5 to 10 us), 0.84x at 1024x10240, 1.84x at 4096x3840 and 3.4x to 3.8x from 4096x10240 up. The floor sits at 12 M elements; above 2**31 elements the flat int32 offsets would wrap, so that falls back too. Flux and qwen-image route only M = 1 modulation GEMMs to nvfp4, so on those the pass was pure launch cost.
# Conflicts: # scripts/build_prequant_checkpoint.py # studio/backend/core/inference/diffusion_prequant.py
danielhanchen
force-pushed
the
pr-10731-ci
branch
from
September 14, 2026 11:00
484a920 to
6bee99f
Compare
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
1 participant
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.
Disposable CI run for unslothai#10731. Do not merge; closed after CI.