Skip to content

Decode NVFP4 block scales once per block, not once per column - #3861

Closed
astorise wants to merge 14 commits into
huggingface:mainfrom
astorise:claude/candle-issue-3860-m5o42c
Closed

astorise wants to merge 14 commits into
huggingface:mainfrom
astorise:claude/candle-issue-3860-m5o42c

Conversation

@astorise

@astorise astorise commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Summary

Nvfp4LinearPacked::forward and dequantize_nvfp4 (candle-transformers/src/quantized_nvfp4.rs)
both decode a block's F8E4M3 weight_scale byte to f32 inside the
per-column loop, instead of once per block_size-wide block:

for (col, val) in row_buf.iter_mut().enumerate() {
    let byte = packed_row[col / 2];
    let nibble = if col % 2 == 0 { byte & 0xf } else { byte >> 4 };
    let block_scale = scale_row[col / self.block_size].to_f32(); // redone every column
    *val = unpack_e2m1(nibble) * block_scale * self.weight_scale_2;
}

With block_size = 16, that's 16x more F8E4M3::to_f32 conversions than
the format requires — the same cost as unpacking a nibble, redone on every
column of the block instead of once per block.

Fixes #3860.

Change

Both loops now walk block-by-block: decode block_scale once per block,
then iterate only the columns inside that block reusing it. Behavior is
unchanged — this is a decode-cost fix, not a numeric one.

Testing

  • cargo test -p candle-transformers --lib quantized_nvfp4 — all 8 tests
    pass, including nvfp4_linear_packed_matches_dense_dequantization
    (bit-exact packed-vs-dense parity) and
    dequantize_nvfp4_matches_manual_computation.
  • cargo clippy -p candle-transformers --lib — no new warnings.
  • cargo fmt -p candle-transformers -- --check — clean.

This only touches the CPU decode path (Nvfp4LinearPacked and the dense
dequantize_nvfp4 fallback); the nvfp4-cuda fused kernel path
(cuda::Nvfp4LinearCuda) is untouched, so no GPU testing is needed for
this change.

astorise and others added 14 commits August 8, 2026 07:11
candle-nvfp4-kernels provided an Nvfp4Linear operator, but nothing in
candle could locate a packed NVFP4 operator in a checkpoint: every
consumer would have needed to parse tensor names, verify shapes, and
extract packed bytes itself (huggingface#3846).

candle-transformers::quantized_nvfp4 adds that weight source:
- detect_backing() inspects a VarBuilder's tensor names to tell dense,
  block-wise FP8, and packed NVFP4 layers apart, so mixed checkpoints
  work layer by layer.
- dequantize_nvfp4() decodes packed E2M1 nibbles and E4M3 block scales
  explicitly (never via a numeric dtype cast), bounded by an explicit
  max_dense_bytes budget so an 8x unpack fails loudly instead of
  blowing past available memory.
- auto_linear() ties detection and construction together: on a CUDA
  device with the new nvfp4-cuda feature it keeps NVFP4 weights packed
  and device-resident via candle-nvfp4-kernels::Nvfp4Linear, otherwise
  it degrades to a dense Linear within the configured memory bound.

Also cherry-picks this fork's three prior commits that introduced
candle-nvfp4-kernels itself, which this weight source depends on.
cuda_capability_signals_are_false_without_cuda_feature asserted
is_available() is always false, but that's only true for the
non-cuda stub; on real CUDA hardware with --features cuda it
correctly returns true, and the unconditional assertion failed
the suite on the GPU runner. Gate it to the case its name and
premise describe.

Found by dispatching CI / cuda against this branch.
cargo clippy --workspace --tests -- -D warnings failed on main:
candle-transformers/src/quantized_nvfp4.rs:469 indexed dequant[row]
inside a for row in 0..out_dim loop instead of iterating directly.
Blocking bug from Tachyon's audit: a real ModelOpt NVFP4 checkpoint
(nvidia/Llama-3.3-70B-Instruct-FP4) has no `weight_packed` tensor --
the packed weight is `.weight` (U8), alongside `.weight_scale`
(F8E4M3, per-block) and `.weight_scale_2` (f32 tensor scale).
`weight_packed` is the compressed-tensors (llm-compressor/vLLM)
convention, not ModelOpt's, despite this module's docs claiming both.
detect_backing required weight_packed, so every ModelOpt NVFP4 layer
fell through to Backing::Dense and candle_nn::linear_b tried to read
the packed U8 bytes as a dense [out, in] tensor in the VarBuilder's
dtype -- a loud failure, but on exactly the checkpoint format the
issue is about.

detect_backing now keys off weight_scale_2 alone (present under both
conventions, absent from dense and block-wise-FP8 layers). A new
nvfp4_weight_tensor_name() picks weight_packed if present, otherwise
falls back to weight, and nvfp4_linear/Nvfp4LinearCuda::new use it
instead of hardcoding weight_packed.

Also from the same audit:
- weight_scale_2 is now read via read_scalar_f32(), which accepts
  either a rank-0 scalar or a rank-1 [1] tensor -- ModelOpt has
  exported it as both across versions, and the old get_with_hints_dtype((), ...)
  call required an exact rank-0 match.
- Documented that input_scale is intentionally never read (W4A16
  doesn't quantize activations), so it doesn't read as an oversight.
- Replaced the hand-rolled e4m3_to_f32 with float8::F8E4M3::to_f32,
  removing a duplicate that could drift from the dependency already
  in the tree.
- Nvfp4LinearCuda::new now reads the packed weight and scale via
  flatten_all()+to_vec1() instead of to_vec2()+concat()/flatten(),
  dropping a transient ~2x-packed-size Vec<Vec<u8>> allocation.

New test nvfp4_linear_reads_the_modelopt_naming_convention exercises
the fix end-to-end against ModelOpt-shaped tensor names.
read_scalar_f32 no longer uses get_unchecked, which VarMap-backed
VarBuilders reject outright ("get_unchecked does not make sense for
VarMap, use get"). It now tries the rank-0 shape and falls back to
[1] with ordinary shape-checked reads. This doesn't make the module
reachable from a ShardedVarBuilder (tensor-parallel loading uses a
different VarBuilderArgs<B> instantiation than the concrete
VarBuilder these functions take) -- that gap is now documented
in-code rather than silent.

Added auto_linear_dispatches_by_detected_backing, which drives
auto_linear directly across all three backings (dense, FP8 block,
NVFP4/ModelOpt) and checks both the returned AutoLinear variant and
forward()'s output shape. The existing
nvfp4_linear_reads_the_modelopt_naming_convention test only ever
called nvfp4_linear directly despite its comment claiming auto_linear
coverage; the comment now says what it actually tests and points at
the new test for the auto_linear path.
auto_linear_dispatches_by_detected_backing built the NVFP4 test row
from a 16-element (BLOCK_SIZE-wide) row instead of an in_dim=4-wide
one, so the packed weight buffer (16 bytes) didn't match its declared
[out_dim, in_dim/2] = [2, 2] shape (4 bytes). This went unnoticed on
candle-core built from sync/upstream-main, where
ShapeWithOneHole::into_shape's blanket impl for concrete shapes never
validated the buffer length; main's candle-core already carries a
local hardening commit that does, and caught it. Use an in_dim-wide
row instead.
…3857)

NVFP4's only non-CUDA execution path dequantized the packed weight to
a dense f32 tensor at load time -- an 8x footprint expansion that OOMs
or refuses to run on CPU hosts and Apple Silicon, where memory is the
actual blocker.

Nvfp4LinearPacked keeps the packed weight and block scales resident in
their checkpoint-native form and decodes each output row on the fly
during forward, instead of pre-expanding the whole matrix. auto_linear
now prefers it whenever the nvfp4-cuda fused kernel isn't available,
with the old dense dequantization kept as an opt-in
(Nvfp4Config::force_dense_fallback) for callers who'd rather trade
memory for throughput. A Metal kernel is still future work; this path
round-trips Metal tensors through host memory in the meantime.
…Linear

Nvfp4LinearPacked pre-decoded weight_scale to Vec<f32> at load time,
so its resident footprint was 1.33x the checkpoint's (4-byte scales
instead of the checkpoint's 1-byte F8E4M3), not the ~1x the docs
claimed. Keep weight_scale as raw F8E4M3 bytes and decode each block
scale in forward's per-row loop instead -- out_dim * n_blocks
conversions per forward, negligible next to the out_dim * batch *
in_dim multiply-accumulates already there. The struct's footprint now
matches the checkpoint's, making the doc comments accurate.

Mark AutoLinear #[non_exhaustive]: this PR just added a variant, and
the module's own docs already flag a Metal-backed variant as expected
future work, so downstream exhaustive matches would break again on
the next backend.
Tachyon's review flagged that Nvfp4LinearPacked computes in f32 and
rounds only the final output, while nvfp4_linear rounds the dense
weight to the VarBuilder's dtype once at load time and compounds that
error into every matmul accumulation -- not a bug, but undocumented
and, since the existing parity test only runs at f32, uncovered.

Add an explicit precision note to Nvfp4LinearPacked's docs, and a test
that builds both paths at reduced precision and checks: packed's
output matches an f32 reference rounded exactly once at the end, while
dense's compounded rounding diverges from that same reference by more
than packed's single rounding step does. Uses F16 rather than BF16,
since candle-core's default CPU matmul only supports F16/F32/F64 --
nvfp4_linear's dense path fails outright at BF16 on CPU today, which
is a pre-existing candle-core gap, not something to route around here.
Tensor::rand is unseeded; an unlucky draw could make packed's single
final-rounding error land as large as dense's compounded one, flaking
the dense_max_abs_diff > packed_max_abs_diff assertion. Confirmed
flaky on main's default feature set (89 passed, this one failed) even
though it passed reliably on the branch this was authored against.
Replaced with a fixed, non-uniform input built from index arithmetic.
…uggingface#3860)

Nvfp4LinearPacked::forward and dequantize_nvfp4 both called
F8E4M3::to_f32 on the block scale inside the per-column loop, redoing
the same conversion block_size times per block instead of once. Hoist
the decode out to the block boundary in both.
@astorise
astorise force-pushed the claude/candle-issue-3860-m5o42c branch from 17a4bf2 to ac480c6 Compare August 8, 2026 07:14
@astorise

Copy link
Copy Markdown
Contributor Author

« Closing — this carries kernel crates that don't exist upstream. Moving it to a separate crate. »

@astorise astorise closed this Aug 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Nvfp4LinearPacked decodes each block scale once per column, not once per block

2 participants