Byte-LUT NF4 dequant in the JIT block stack - #35
Merged
Conversation
The nf4_lin dequant chain (shift/and to int64, stack, flatten, fp32 gather, scale, cast, cat) generated ~70 bytes of memory traffic per half-byte parameter: ~1.4 TB per denoise step over the 19B weights, matching the measured ~4.4 s fixed per-step cost. Replace it with a single embedding gather through a precomputed [256, 2] byte LUT (hi nibble, lo nibble) built in the compute dtype, absmax cast down before the scale so nothing promotes to float32, no chunk loop, no cat. Isolated TorchScript benchmark (100M params, CUDA): 23.4 -> 4.1 ms, 5.7x; projected ~4.4 -> ~0.8 s of dequant per step. On float32 (CPU, CI) the math is identical to the old path; on bf16 the scale multiply rounds once more (~0.4%, well under NF4 quantization noise). End-to-end GPU parity and render timing to follow.
Contributor
Author
|
Post-merge validation (the "GPU parity and render timing to follow" from the first commit): Render timing (768x512x49, NF4, seed 42, RTX 5060 Ti): denoise steps 6.4–6.5 s → 3.2–3.3 s (1.97x); render wall 89–96 s → 73 s. Remaining wall is transformer onload (~15 s), denoise (~26 s), video decode (12.3 s). Real-weight dequant parity (old chain vs byte LUT, both vs fp32-exact, on artifact weights, bf16 CUDA): the LUT path's max deviation from fp32 truth is 1–2x the old path's own bf16-cast rounding (e.g. attn1.to_k: old 0.094, new 0.147, in units of tensor RMS on the worst element), orders of magnitude below NF4 quantization error. Seed-42 output video archived for QC. |
This was referenced Jul 20, 2026
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.
The LTX-2.3 JIT block stack dequantizes NF4 weights through a precomputed [256, 2] byte lookup table (one embedding gather in the compute dtype) instead of the int64 shift/stack/gather chain, cutting per-step dequant memory traffic ~6x. Isolated benchmark: 5.7x; the dequant was the measured per-step wall at ~4.4 s of a 6.5 s step.
Includes a merge of main to pick up the Qwen3 mask dtype fix (#34) and the corresponding 0.1.0.9 re-bump.