feat(loader): native U8 NVFP4 loading for pre-quantized checkpoints

[Sujimoshi]

Summary

Third in a planned split of #218 (DFlash speculative decoding) into smaller, independently reviewable PRs (first two: #255 merged, #256 closed as unneeded).

Some pre-quantized checkpoints (Standard NVFP4 convention, e.g. sakamakismile) ship weights already packed as U8 NVFP4 on disk at the .weight key (rather than .weight_packed, the compressed-tensors convention). This adds a direct-load path that skips the BF16-dequant-then-requantize roundtrip:

• FullAttention Q/K/V/O (TP-sharded path): detect UInt8 dtype, load via quantized_auto(Nvfp4Variant::Standard) instead of dense_auto (which bails on UInt8). TP sharding of pre-quantized NVFP4 isn't supported yet (tp_size=1 only).
• LinearAttention (SSM/GDN): new native_nvfp4 fast path — when in_proj_qkv/in_proj_z/out_proj are all UInt8, load+concat+transpose them as QuantizedWeight directly. load_ssm_proj also gained a UInt8 branch for partial-U8 checkpoints that fall through to the legacy path. in_proj_a/in_proj_b now route through load_ssm_proj instead of raw dense() so a U8-packed A/B doesn't get read as BF16 garbage.
• quant_helpers.rs: quantized_v2 now sets weight_scale_2_vec (a field added to QuantizedWeight since this code was originally written).

Note on provenance: re-derived against current main rather than patched mechanically. Main has since added a native FP8 GDN fast path (fixing a ~7pt BFCL-ST regression) and an ATLAS_GDN_BF16_WEIGHTS opt-in path — both sit ahead of this new native-NVFP4 branch in priority and are untouched. Common SSM loads (A/B, conv1d, A_log, dt_bias, norm, ba_dense, qkvz_size) are hoisted once above all branches so the new fast path can share them.

Test plan

• cargo check -p spark-model (ATLAS_SKIP_BUILD=1, CUDA stubbed) — clean
• cargo fmt --check -p spark-model — clean
• cargo clippy -p spark-model --tests (ATLAS_SKIP_BUILD=1) — no warnings
• Runtime verification with an actual pre-quantized Standard-NVFP4 checkpoint on GB10 — see comment below

🤖 Generated with Claude Code

[Sujimoshi]

Runtime verification on GB10 — done ✅

Checkpoint: Qwen3.6-27B-NVFP4 (local, modelopt-produced, confirmed Standard NVFP4 convention: dtype=U8 at .weight + separate weight_scale/weight_scale_2/input_scale, for both GDN/SSM projections and FullAttention Q/K/V/O — checked directly against the safetensors headers, not just key names).

Build: cargo build --release -p spark-server with ATLAS_TARGET_MODEL=qwen3.6-27b ATLAS_TARGET_QUANT=nvfp4 (needed explicitly — default kernel target doesn't match this checkpoint's hidden_size/model_type) and LIBRARY_PATH=/tmp/nccl-stubs (no system NCCL on this box). Both environment gaps, unrelated to this PR's code.

Server: spark serve --model-from-path .../Qwen3.6-27B-NVFP4 --kv-cache-dtype nvfp4 --gpu-memory-utilization 0.75 --max-batch-size 1 --max-seq-len 4096, TP=1. Weights loaded in ~7s, 19.15 GB.

Native-path confirmation:

• GDN: all 48 SSM layers logged SSM[layers.N] native NVFP4 GDN: qkvz+out_proj loaded pre-quantized (U8-packed on disk; no BF16 dequant/requant roundtrip) — exactly 48/48, no fallbacks.
• FullAttention Q/K/V/O: fast path confirmed in code (qwen35_dense.rs:291-309) and indirectly in logs — the only dense_auto/quantize_to_nvfp4 log hit in the entire run was the LM head (n=248077,k=5120, i.e. vocab×hidden), not any attention projection. If attention had fallen back to BF16-dequant, each of the 16 attention layers would've logged similarly — none did.

Smoke test: two prompts via /v1/chat/completions ("capital of France", "haiku about a GPU") — both coherent, correct, no repetition/degradation. ~14 tok/s single-batch (expected, no speculative decoding active), TTFT ~128ms.

Finding (not a bug, logging gap): the GDN branch has an explicit tracing::info! confirming the native path; the attention branch doesn't. Recommend adding an analogous log line in a follow-up so the fast path is directly log-visible instead of inferred from absence of fallback logs. Not fixed here since this PR was scoped to testing only.

Server was shut down cleanly after the test, no leftover process.

[rrstesiak]

1. qwen35_dense.rs - documented constraint with no code behind it.
The FullAttention U8 fast path comment says "TP sharding of pre-quantized NVFP4
isn't supported yet (tp_size=1 only)" - but there is no tp_size check. At
tp_size>1 every rank silently loads the full unsharded weight and serves wrong
results. This needs && tp_size == 1 (falling through to the dequant path) or
a hard bail.

2. Same path - NULL dense weight handed to a caller that frees it.
The early return produces DenseWeight { weight: NULL }, and the call site
unconditionally gpu.free()s the dense intermediates after quantization.
Please confirm gpu.free(NULL) is a no-op on every backend this compiles for,
or skip the free for the U8 path.

3. quantized.rs::concat_rows - bit-exact f32 equality as a load gate.
ensure!(self.weight_scale_2 == other.weight_scale_2) means any checkpoint
whose quantizer emitted independent per-tensor scales for in_proj_qkv vs
in_proj_z fails at load. It fails loud (good), but this is a checkpoint-
convention assumption, not an invariant - worth a comment naming which
quantizer guarantees it, and an error message saying "re-quantize with X".

4. loaders_fp8.rs::quantize_to_nvfp4_per_row - dead code, wrong comment.
#[allow(dead_code)], no caller in this PR, and the comment says "use max for
compatibility diagnostics" while the code reads row 0. Belongs in the PR that
calls it, with the comment matching the code.

[Sujimoshi]

Thanks for the thorough review — going through each:

1. TP-sharding gap on the U8 fast path — confirmed, real bug, you're right. There's no tp_size guard at all on that branch; at tp_size>1 it'd silently load the full unsharded weight per rank. Fixing now (loud bail on tp_size>1 rather than a silent wrong result, since sharding pre-quantized NVFP4 isn't implemented yet).

2. gpu.free(NULL) — checked all three backends (cuda_backend/gpu_impl.rs, metal_backend.rs, gpu/mock.rs): free() explicitly no-ops on a null pointer in all of them (CUDA and Metal check ptr.is_null() up front; the mock just no-ops removing a nonexistent key). So this is safe as-is, no code change needed there.

3. Bit-exact weight_scale_2 equality gate — agreed it's a checkpoint-convention assumption rather than a hard invariant. Adding a comment naming the quantizer this holds for and a clearer error message.

4. Dead quantize_to_nvfp4_per_row with a stale comment — confirmed it's new in this PR with no caller yet, and the comment contradicts the code (reads row 0, not max). Removing it from this PR — it belongs in whichever future PR actually wires up the per-row-scale GEMV path.

Pushing a fix for 1/3/4 shortly.

[Sujimoshi]

Pushed fixes for 1/3/4 in `b3469ae`:

1. TP-sharding guard: the U8 fast path now bails loudly (`anyhow::ensure!`) if `tp_size > 1`, instead of silently loading the full unsharded weight per rank.
2. `concat_rows` comment/error message: clarified that the bit-exact `weight_scale_2` equality is valid specifically for ModelOpt/Standard NVFP4 (single global per-tensor scalar), not a general invariant across all conventions.
3. Dead code removed: `quantize_to_nvfp4_per_row` deleted — confirmed no callers anywhere in the codebase and new to this PR; will come back in whichever future PR actually wires up the per-row-scale GEMV path.

`cargo check`/`fmt --check`/`clippy --tests -p spark-model` all clean after the changes. Thanks again for catching the TP bug — that one was real.

[pragmaxim]

@Sujimoshi I just tested it with :

./target/release/spark serve morosystems/ThinkingCap-Qwen3.6-27B-NVFP4 \
      --port 8888 \
      --host 0.0.0.0 \
      --max-seq-len 65536 \
      --kv-cache-dtype bf16 \
      --gpu-memory-utilization 0.9 \
      --scheduling-policy slai \
      --enable-prefix-caching \
      --tool-call-parser qwen3_coder

And it fixed the error 👍

[SeedSource]

[BETA] Automated PR review — SeedSource AI pipeline

Advisory only, not a merge gate. AI-generated (internal risk-triage, Atlas serve gate, GLM-5.2 review), attributed per Atlas convention.

Triage

• Risk: high - serve-test-required (triage advisory): True - path-heuristic (authority): True
• Rationale: Adds NVFP4 pre-quantized U8 handling and enforces tp_size=1 to prevent silent TP sharding failures.
• Serving-path files: crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs, crates/spark-model/src/tp_shard/quant_shard.rs, crates/spark-model/src/weight_loader/qwen35_dense.rs, crates/spark-model/src/weight_loader/step3p7.rs, crates/spark-model/src/weight_loader/step3p7/load_layers.rs, crates/spark-model/src/weight_map/loaders_fp8.rs

Serve gate

• Build: 19s - Qwen3.6-35B-A3B-FP8 on GB10 node1
• Top-k KL vs main: 0.0002 (no shift vs main) - decode delta -0.1%
• Coherence probe: coherence=pass ping=10/10 (note: our test_coherence.py is flaky here; KL is the reliable serve signal)

GLM-5.2 code review

Correctness concerns:

1. SSM out_proj field mismatch: In the native NVFP4 path, SsmWeights.out_proj is assigned out_proj_nvfp4 (a QuantizedWeight), but the legacy path assigns out_proj_dense (a DenseWeight). If SsmWeights.out_proj is typed DenseWeight, this won't compile; if it's a union/enum, the downstream consumer must handle both — not visible in this diff. Verify the type and that Qwen3SsmLayer::new_sequential + predequant_for_prefill correctly handle the quantized out_proj path.

2. weight_scale_2_vec propagation: Added to struct construction in two places but no sharding logic for it. For tp_size=1 this is fine (pass-through), but the field is new — confirm it's initialized everywhere else or defaults to NULL.

3. in_proj_a/in_proj_b now routed through load_ssm_proj with nv/h dims instead of raw dense(). If A/B were previously BF16 and are now potentially U8-dequantized, confirm interleave_ba still receives BF16. The comment claims this is correct but no assertion enforces it.

4. quantized_auto for SSM projections — no tp_size guard here unlike the attention path. If tp_size > 1 reaches this code, quantized_auto may load full unsharded weights silently. Add the same anyhow::ensure!(tp_size == 1) guard.

Risk: Medium. The SSM native path is substantial new code with no visible tests. Partial-U8 checkpoints fall through to legacy path — verify load_ssm_proj's new U8 branch produces correct shapes for concat.

Quality: Comments are excessive (multi-paragraph essays). Trim to essential context.

_{Pipeline: build+serve on GB10, GLM-5.2 via NVIDIA NIM, internal risk-triage classifier. External diff treated as untrusted input. Public repo.}