finegrained-kernels: decode losslessly compressed NVFP4 scales (CSF)

[IlyasMoutawwakil]

Adds GPU decoding for CSF-compressed NVFP4 weight scales, as used by Qwen3.8-Flash-Next-FP4-Local. This enables callers to read those checkpoint tensors and choose between keeping scales expanded or saving GPU memory by decoding them when needed.

NVFP4 stores 4-bit weights plus an 8-bit scale for every 16 weights. CSF compresses the scale bytes losslessly: it stores a base byte per row, packs most scales as 4-bit offsets from that base, and keeps values outside that range in an exception stream. Decoding reconstructs the original FP8 scale bytes exactly, so this compression adds no quantization error. The packed FP4 weights stay unchanged.

The benefit is smaller scale storage. Keeping that saving during inference costs additional decoding latency with this implementation. Callers have two choices:

Strategy	GPU memory	Inference cost
Decode once when loading; release compressed buffers	Ordinary dense scale storage	Ordinary NVFP4 GEMM; no repeated decoding
Keep scales compressed; decode before each GEMM into shared workspace	Compressed scales for resident layers, plus one reusable dense workspace	Additional GPU decoding before each GEMM

For the measured 512-expert Qwen gate/down projection stacks, compressed scales occupy 28.68 / 27.99 MB instead of 52.43 MB each, about 45–47% less scale storage. These percentages apply to scales, not total model memory. Across all 73,728 compressed scale planes in the actual checkpoint (48 layers × 512 experts × gate/up/down), CSF reduces scales from 7.03 GiB to 3.69 GiB. The complete checkpoint tensor payload is 95.12 GiB with CSF versus a calculated 98.46 GiB with the same scales expanded: 3.34 GiB saved, or 3.39% of total tensor storage. These counts come from all relevant safetensors headers; runtime workspace and other allocations are additional. Staging also needs a 52.43 MB dense workspace that the caller can reuse across projections/layers. For a single projection stack, compressed buffers plus workspace exceed dense storage; the memory benefit comes from sharing that workspace across multiple resident stacks.

The left panel compares stored scales. The right panel includes the shared workspace and illustrates how total scale memory changes as more gate/down pairs remain resident. This is an extrapolation of the measured tensor sizes, not a measured whole-model memory footprint.

Fresh B200 measurements use the optimized decoder and measure decode + GEMM together, comparing with the same ordinary NVFP4 GEMM using dense scales. Warm and L2-flushed CUDA-graph replay are shown separately on linear axes. Across the four cases, staging adds 16.5–28.8 µs, or approximately 9–32% latency. This is a capacity tradeoff, with no demonstrated inference speedup. The decoder now processes contiguous pairs instead of padded row tiles; the latency plot includes this optimization.

These are kernel measurements using actual first-expert scale streams repeated across 512 experts, random packed weights, and the same existing 100-trial GEMM tuning winners, with no retuning during measurement. They are not end-to-end model throughput. Both strategies preserve the original scale bytes exactly. With 8 of 512 experts active, reusing routing metadata reduces gate decoding from 20.07 to 5.65 µs and complete decode + batched GEMM from 29.61 to 13.20 µs (warm graphs; the same GEMM with resident dense scales takes 7.07 µs). Cold pipeline timings are 36.93 → 22.82 µs, versus 16.45 µs dense. These sparse results are separate from the all-expert plot above; sparse down/prefill timings were not collected because their cached tuning winners were unavailable.

The change adds prepare_nvfp4_csf_scales and decompress_nvfp4_csf_scales in the existing loading/scales.py. Preparation preserves the input device; decoding runs on the GPU and accepts reusable workspace. GEMM implementations and dispatch stay unchanged. The caller owns the staging/workspace policy. Pass the existing grouped routing expert_start tensor to decode only experts with routed tokens; both the fixed stream and exception pass skip inactive experts, leaving their workspace planes untouched. No selection buffer or CPU synchronization is added. Without routing metadata, the decoder expands the full supplied stack. This PR supplies the helpers, not automatic model-loader integration.

Validation: 11 passed, 1 expected routing-only skip in the shared end-to-end scenarios on B200, covering dense/batched/grouped GEMMs, exception reconstruction, reusable odd-offset workspace, and fullgraph decode + gated GEMM + requantization. Decoder output was also checked against actual checkpoint streams. Ruff and diff checks passed.

• Linked to a discussed issue (requested directly by a maintainer; no existing issue)
• Tested locally
• New functionality has test coverage
• Mostly or completely generated with an LLM

[github-actions]

This PR changes source code for the following kernel(s): finegrained-kernels.

If you have the kernel-bot build permissions, then trigger /kernel-bot build finegrained-kernels yourself. Otherwise, a maintainer will review the PR and will take an appropriate action. Thanks for your patience.