Kernels:

phanerozoic
/

fp4-train

Kernel card Files Files and versions

xet

Community

phanerozoic commited on Jul 20

Commit

75066be

verified ·

1 Parent(s): 2a71500

Card: standardized form with hero

Browse files

Files changed (2) hide show

CARD.md +63 -78
README.md +63 -78

CARD.md CHANGED Viewed

@@ -5,15 +5,27 @@ license: apache-2.0
 # fp4-train
-fp4-train is an NVFP4 quantizer with stochastic rounding and error feedback, the
-primitives for stable 4-bit training, loadable through `kernels`. Low-bit
-training is limited less by the matmul than by the quantizer: round-to-nearest is
-biased, so the quantization error correlates with the signal and the optimizer
-drifts. This kernel quantizes to NVFP4 (e2m1 values with an e4m3 scale per
-16-element block) with stochastic rounding, so the quantization is unbiased, and
-carries a per-element error-feedback residual so sub-ULP information accumulates
-instead of vanishing. The packed output is in the two-codes-per-byte layout the
-Blackwell FP4 tensor cores consume.
 ## Usage
@@ -23,28 +35,19 @@ from kernels import get_kernel
 fp4 = get_kernel("phanerozoic/fp4-train", version=1, trust_remote_code=True)
-# NVFP4-stored-weight training: the weight lives in 4-bit and is updated in place.
-# Error feedback carries the sub-ULP part of each update, so round-to-nearest no
-# longer stalls; this is the regime where SR/EF matter.
 ef = fp4.ErrorFeedback(stochastic=False)                      # RTN + error feedback
-W = W0.clone()                                                # the NVFP4 weight (bf16 values)
 for step in range(num_steps):
-    grad = compute_grad(W)                                    # your backward
-    W = ef.qdq("layer0.weight", W - lr * grad, step=step)     # apply the update in NVFP4
-# stochastic=True instead gives unbiased stochastic rounding (a noisier estimator)
 # quantization primitives
-packed, scale = fp4.quantize(W, stochastic=True, seed=0, step=global_step)  # NVFP4 + e4m3
 Wq = fp4.dequantize(packed, scale, *W.shape)
-# quantized-operand linear (full-precision master weights kept by the caller)
-x = torch.randn(8, 512, 2048, device="cuda")
-Wm = torch.randn(6144, 2048, device="cuda", requires_grad=True)
-y = fp4.fp4_linear(x, Wm, stochastic=True, step=global_step)  # NVFP4 operands, STE grad
 # native NVFP4 tensor-core matmul on Blackwell (a @ b^T, K a multiple of 16)
-a = torch.randn(2048, 4096, device="cuda"); b = torch.randn(4096, 4096, device="cuda")
-y = fp4.fp4_mm(a, b)                                          # hardware FP4 path, bf16 out
 ```
 `version` selects the release branch; `trust_remote_code` is required by
@@ -57,79 +60,61 @@ y = fp4.fp4_mm(a, b)                                          # hardware FP4 pat
 | `quantize(x, stochastic, residual, seed, step)` | NVFP4 quantize; error feedback if `residual` is a `[M,K]` f32 tensor (updated in place) |
 | `dequantize(packed, scale, M, K)` | NVFP4 -> bf16 |
 | `qdq(x, stochastic, residual, seed, step)` | quantize then dequantize (the value the tensor cores use) |
-| `fp4_linear(x, weight, stochastic, seed, step)` | training-ready quantized-operand linear (straight-through gradient) |
 | `fp4_matmul(a, b, stochastic, seed, step)` | `a @ b^T` FP4-operand product, dequant + bf16 accumulate (portable, cc 8.0+) |
 | `fp4_mm(a, b, stochastic, seed, step)` | `a @ b^T` on the Blackwell NVFP4 tensor cores (native `_scaled_mm`, ~10x faster) |
 | `to_fp4_operands(x, stochastic, seed, step)` | packed fp4 + e4m3 scale as the torch FP4 dtypes for `_scaled_mm` |
 | `ErrorFeedback` | per-tensor residual state for weight quantization |
-## How it works
-Each 16-element block gets a scale of `block_absmax / 6` encoded to e4m3 (so the
-largest element maps near the e2m1 maximum of 6). Every element is divided by the
-decoded scale and encoded to the e2m1 grid `{0, .5, 1, 1.5, 2, 3, 4, 6}` with a
-sign bit. Rounding is stochastic: for a value between two grid points, the higher
-is chosen with probability equal to the fractional distance, using a counter-based
-(seed, step, global-index) RNG, so the expected quantized value equals the input.
-With error feedback, the per-element residual `x - dequant(quantize(x))` is
-carried and added into the next quantization, so the time-averaged quantized
-value equals the true value even under round-to-nearest. Codes are packed two per
-byte; the block scales are e4m3, the layout NVFP4 tensor cores read.
 ## Correctness
 Measured on L4, H200, and RTX PRO 6000 through `get_kernel`:
-- **Exact encoding.** The quantize/dequantize matches an independent reference
-  implementation to 0 (`max abs diff = 0`).
-- **Stochastic rounding is unbiased.** For a value 2.3 sitting between the grid
-  points 2 and 3, the stochastic-rounding mean is 2.300 against a true 2.3, where
-  round-to-nearest is biased to 2.000.
-- **Error feedback cancels the drift.** With round-to-nearest plus error
-  feedback the time-averaged quantized value is 2.297, tracking the true 2.3,
-  where plain round-to-nearest stays stuck at 2.000.
-- **Round-to-nearest stalls low-bit training; stochastic rounding and error
-  feedback do not.** With the weights *stored* in NVFP4 and updated in place, so
-  each step's update is smaller than a quantization step, round-to-nearest
-  discards every update and training stalls, while stochastic rounding (unbiased)
-  and error feedback (residual-carrying) apply the updates and recover
-  near-bf16 loss. A pure accumulation of a sub-ULP increment reaches a true 1.50
-  under SR (1.50) and EF (1.50) while RTN stays at the starting 0.50. A linear
-  model trained with NVFP4-stored weights reaches MSE 1.0 (EF) and 13 (SR) against
-  a bf16 0.0, where RTN sits at 77; a small MLP reaches loss 0.159 (EF) and 0.171
-  (SR) against a bf16 0.158, where RTN stalls at 1.03. Error feedback recovers
-  essentially the full-precision loss; stochastic rounding, being a noisier
-  estimator, trails it but still beats RTN by a wide margin.
-## Scope
-The verified contribution is the quantizer that makes 4-bit training stable:
-correct NVFP4 encoding, unbiased stochastic rounding, and error feedback that
-preserves sub-ULP information. Their value shows up when the low precision is in
-the training state itself, meaning weights stored and updated in NVFP4, where the
-per-step update is below a quantization step: there round-to-nearest stalls and
-stochastic rounding or error feedback recover near-bf16 loss, as measured above.
-When instead only the forward matmul operands are quantized and a full-precision
-master weight is kept, the choice of rounding is second order, since the master
-carries the sub-ULP information; that regime does not need this kernel's rounding
-modes and is not where the gain is. Two matmuls are bundled: `fp4_matmul` takes a
-numerically exact FP4 path (dequantize, bf16 accumulate) that is portable to every
-compute-capability-8.0+ GPU, and `fp4_mm` runs the operands on the Blackwell NVFP4
-tensor cores by swizzling the e4m3 block scales into the layout `torch._scaled_mm`
-consumes. On an RTX PRO 6000 the hardware path returns a bit-identical result
-(relative error 0) and is about 10-14x faster than the dequantize path and 1.5-1.7x
-faster than a bf16 matmul.
 ## Requirements and limits
-- NVIDIA GPU with compute capability 8.0+ (e4m3 scale via software conversion).
 - `K` a multiple of 16; float32 or bfloat16 input.
 ## References
 NVFP4 and OCP microscaling (MXFP) formats; Gupta et al., "Deep Learning with
-Limited Numerical Precision" (stochastic rounding, 2015); Seide et al., "1-bit
-SGD" (error feedback, 2014); low-precision training with unbiased quantization.
 ## License

 # fp4-train
+fp4-train is an NVFP4 quantizer with stochastic rounding and error feedback,
+the primitives for stable 4-bit training, loadable through `kernels`. The
+reference baseline is an independent NVFP4 encoding implementation, matched
+to zero error, with fp32 training as the convergence reference.
+Training with weights stored in 4-bit fails under ordinary rounding: once the
+per-step update is smaller than a quantization step, round-to-nearest
+discards every update and the model stops learning. This kernel quantizes to
+NVFP4 (e2m1 values with an e4m3 scale per 16-element block) with stochastic
+rounding, so quantization is unbiased and sub-step updates accumulate in
+expectation, and it carries a per-element error-feedback residual so the
+dropped fraction of every update is banked instead of lost. The packed output
+is the layout Blackwell FP4 tensor cores consume.
+![Four training curves: round-to-nearest flatlines while stochastic rounding and error feedback descend toward the fp32 reference, above an e2m1 grid where a value hops between neighbors](https://huggingface.co/kernels/phanerozoic/fp4-train/resolve/main/media/hero.gif)
+*Weights stored in NVFP4 and updated in place, 500 steps of sub-step updates:
+round-to-nearest never leaves its starting loss (6.56); stochastic rounding
+reaches 1.80, error feedback 0.77, against the fp32 reference at 0.02. Below,
+quantizing 2.3 on the e2m1 grid hops between 2 and 3 with the running mean at
+2.26.*
 ## Usage
 fp4 = get_kernel("phanerozoic/fp4-train", version=1, trust_remote_code=True)
+# NVFP4-stored-weight training: the weight lives in 4-bit, updated in place
 ef = fp4.ErrorFeedback(stochastic=False)                      # RTN + error feedback
+W = W0.clone()
 for step in range(num_steps):
+    grad = compute_grad(W)
+    W = ef.qdq("layer0.weight", W - lr * grad, step=step)
 # quantization primitives
+packed, scale = fp4.quantize(W, stochastic=True, seed=0, step=global_step)
 Wq = fp4.dequantize(packed, scale, *W.shape)
 # native NVFP4 tensor-core matmul on Blackwell (a @ b^T, K a multiple of 16)
+y = fp4.fp4_mm(a, b)
 ```
 `version` selects the release branch; `trust_remote_code` is required by
 | `quantize(x, stochastic, residual, seed, step)` | NVFP4 quantize; error feedback if `residual` is a `[M,K]` f32 tensor (updated in place) |
 | `dequantize(packed, scale, M, K)` | NVFP4 -> bf16 |
 | `qdq(x, stochastic, residual, seed, step)` | quantize then dequantize (the value the tensor cores use) |
+| `fp4_linear(x, weight, stochastic, seed, step)` | quantized-operand linear (straight-through gradient) |
 | `fp4_matmul(a, b, stochastic, seed, step)` | `a @ b^T` FP4-operand product, dequant + bf16 accumulate (portable, cc 8.0+) |
 | `fp4_mm(a, b, stochastic, seed, step)` | `a @ b^T` on the Blackwell NVFP4 tensor cores (native `_scaled_mm`, ~10x faster) |
 | `to_fp4_operands(x, stochastic, seed, step)` | packed fp4 + e4m3 scale as the torch FP4 dtypes for `_scaled_mm` |
 | `ErrorFeedback` | per-tensor residual state for weight quantization |
+## Method
+Each 16-element block gets a scale of `block_absmax / 6` encoded to e4m3, so
+the largest element maps near the e2m1 maximum. Every element is divided by
+the decoded scale and encoded to the e2m1 grid `{0, .5, 1, 1.5, 2, 3, 4, 6}`
+with a sign bit. Rounding is stochastic: for a value between two grid points,
+the higher is chosen with probability equal to the fractional distance, using
+a counter-based `(seed, step, global-index)` RNG, so the expected quantized
+value equals the input. With error feedback, the residual
+`x - dequant(quantize(x))` is carried and added into the next quantization,
+so the time-averaged value equals the true value even under round-to-nearest.
+Codes pack two per byte; block scales are e4m3, the layout NVFP4 tensor cores
+read.
 ## Correctness
 Measured on L4, H200, and RTX PRO 6000 through `get_kernel`:
+- Exact encoding: quantize/dequantize matches an independent reference to 0
+  (`max abs diff = 0`).
+- Stochastic rounding is unbiased: for 2.3 between grid points 2 and 3, the
+  SR mean is 2.300 where round-to-nearest is biased to 2.000.
+- Error feedback cancels the drift: RTN plus error feedback time-averages to
+  2.297.
+- With NVFP4-stored weights and sub-step updates, RTN stalls at the starting
+  loss while SR and EF recover near-full-precision loss: a linear model
+  reaches MSE 1.0 (EF) and 13 (SR) against bf16 0.0 with RTN at 77; a small
+  MLP reaches 0.159 (EF) and 0.171 (SR) against bf16 0.158 with RTN at 1.03.
+- On an RTX PRO 6000 the hardware `fp4_mm` path returns a bit-identical
+  result to the dequantize path (relative error 0) and is about 10-14x
+  faster, and 1.5-1.7x faster than a bf16 matmul.
 ## Requirements and limits
+- NVIDIA GPU with compute capability 8.0+ (e4m3 scale via software
+  conversion); `fp4_mm` requires Blackwell.
 - `K` a multiple of 16; float32 or bfloat16 input.
+- The rounding modes matter when the low precision is in the training state
+  itself (weights stored and updated in NVFP4). With a full-precision master
+  weight, the master carries the sub-ULP information and the choice of
+  rounding is second order.
+- Blocks whose absmax falls below the e4m3 scale range quantize to zero; keep
+  stored-weight scales inside the representable range.
 ## References
 NVFP4 and OCP microscaling (MXFP) formats; Gupta et al., "Deep Learning with
+Limited Numerical Precision" (stochastic rounding, 2015); Seide et al.,
+"1-bit SGD" (error feedback, 2014).
 ## License

README.md CHANGED Viewed

@@ -5,15 +5,27 @@ license: apache-2.0
 # fp4-train
-fp4-train is an NVFP4 quantizer with stochastic rounding and error feedback, the
-primitives for stable 4-bit training, loadable through `kernels`. Low-bit
-training is limited less by the matmul than by the quantizer: round-to-nearest is
-biased, so the quantization error correlates with the signal and the optimizer
-drifts. This kernel quantizes to NVFP4 (e2m1 values with an e4m3 scale per
-16-element block) with stochastic rounding, so the quantization is unbiased, and
-carries a per-element error-feedback residual so sub-ULP information accumulates
-instead of vanishing. The packed output is in the two-codes-per-byte layout the
-Blackwell FP4 tensor cores consume.
 ## Usage
@@ -23,28 +35,19 @@ from kernels import get_kernel
 fp4 = get_kernel("phanerozoic/fp4-train", version=1, trust_remote_code=True)
-# NVFP4-stored-weight training: the weight lives in 4-bit and is updated in place.
-# Error feedback carries the sub-ULP part of each update, so round-to-nearest no
-# longer stalls; this is the regime where SR/EF matter.
 ef = fp4.ErrorFeedback(stochastic=False)                      # RTN + error feedback
-W = W0.clone()                                                # the NVFP4 weight (bf16 values)
 for step in range(num_steps):
-    grad = compute_grad(W)                                    # your backward
-    W = ef.qdq("layer0.weight", W - lr * grad, step=step)     # apply the update in NVFP4
-# stochastic=True instead gives unbiased stochastic rounding (a noisier estimator)
 # quantization primitives
-packed, scale = fp4.quantize(W, stochastic=True, seed=0, step=global_step)  # NVFP4 + e4m3
 Wq = fp4.dequantize(packed, scale, *W.shape)
-# quantized-operand linear (full-precision master weights kept by the caller)
-x = torch.randn(8, 512, 2048, device="cuda")
-Wm = torch.randn(6144, 2048, device="cuda", requires_grad=True)
-y = fp4.fp4_linear(x, Wm, stochastic=True, step=global_step)  # NVFP4 operands, STE grad
 # native NVFP4 tensor-core matmul on Blackwell (a @ b^T, K a multiple of 16)
-a = torch.randn(2048, 4096, device="cuda"); b = torch.randn(4096, 4096, device="cuda")
-y = fp4.fp4_mm(a, b)                                          # hardware FP4 path, bf16 out
 ```
 `version` selects the release branch; `trust_remote_code` is required by
@@ -57,79 +60,61 @@ y = fp4.fp4_mm(a, b)                                          # hardware FP4 pat
 | `quantize(x, stochastic, residual, seed, step)` | NVFP4 quantize; error feedback if `residual` is a `[M,K]` f32 tensor (updated in place) |
 | `dequantize(packed, scale, M, K)` | NVFP4 -> bf16 |
 | `qdq(x, stochastic, residual, seed, step)` | quantize then dequantize (the value the tensor cores use) |
-| `fp4_linear(x, weight, stochastic, seed, step)` | training-ready quantized-operand linear (straight-through gradient) |
 | `fp4_matmul(a, b, stochastic, seed, step)` | `a @ b^T` FP4-operand product, dequant + bf16 accumulate (portable, cc 8.0+) |
 | `fp4_mm(a, b, stochastic, seed, step)` | `a @ b^T` on the Blackwell NVFP4 tensor cores (native `_scaled_mm`, ~10x faster) |
 | `to_fp4_operands(x, stochastic, seed, step)` | packed fp4 + e4m3 scale as the torch FP4 dtypes for `_scaled_mm` |
 | `ErrorFeedback` | per-tensor residual state for weight quantization |
-## How it works
-Each 16-element block gets a scale of `block_absmax / 6` encoded to e4m3 (so the
-largest element maps near the e2m1 maximum of 6). Every element is divided by the
-decoded scale and encoded to the e2m1 grid `{0, .5, 1, 1.5, 2, 3, 4, 6}` with a
-sign bit. Rounding is stochastic: for a value between two grid points, the higher
-is chosen with probability equal to the fractional distance, using a counter-based
-(seed, step, global-index) RNG, so the expected quantized value equals the input.
-With error feedback, the per-element residual `x - dequant(quantize(x))` is
-carried and added into the next quantization, so the time-averaged quantized
-value equals the true value even under round-to-nearest. Codes are packed two per
-byte; the block scales are e4m3, the layout NVFP4 tensor cores read.
 ## Correctness
 Measured on L4, H200, and RTX PRO 6000 through `get_kernel`:
-- **Exact encoding.** The quantize/dequantize matches an independent reference
-  implementation to 0 (`max abs diff = 0`).
-- **Stochastic rounding is unbiased.** For a value 2.3 sitting between the grid
-  points 2 and 3, the stochastic-rounding mean is 2.300 against a true 2.3, where
-  round-to-nearest is biased to 2.000.
-- **Error feedback cancels the drift.** With round-to-nearest plus error
-  feedback the time-averaged quantized value is 2.297, tracking the true 2.3,
-  where plain round-to-nearest stays stuck at 2.000.
-- **Round-to-nearest stalls low-bit training; stochastic rounding and error
-  feedback do not.** With the weights *stored* in NVFP4 and updated in place, so
-  each step's update is smaller than a quantization step, round-to-nearest
-  discards every update and training stalls, while stochastic rounding (unbiased)
-  and error feedback (residual-carrying) apply the updates and recover
-  near-bf16 loss. A pure accumulation of a sub-ULP increment reaches a true 1.50
-  under SR (1.50) and EF (1.50) while RTN stays at the starting 0.50. A linear
-  model trained with NVFP4-stored weights reaches MSE 1.0 (EF) and 13 (SR) against
-  a bf16 0.0, where RTN sits at 77; a small MLP reaches loss 0.159 (EF) and 0.171
-  (SR) against a bf16 0.158, where RTN stalls at 1.03. Error feedback recovers
-  essentially the full-precision loss; stochastic rounding, being a noisier
-  estimator, trails it but still beats RTN by a wide margin.
-## Scope
-The verified contribution is the quantizer that makes 4-bit training stable:
-correct NVFP4 encoding, unbiased stochastic rounding, and error feedback that
-preserves sub-ULP information. Their value shows up when the low precision is in
-the training state itself, meaning weights stored and updated in NVFP4, where the
-per-step update is below a quantization step: there round-to-nearest stalls and
-stochastic rounding or error feedback recover near-bf16 loss, as measured above.
-When instead only the forward matmul operands are quantized and a full-precision
-master weight is kept, the choice of rounding is second order, since the master
-carries the sub-ULP information; that regime does not need this kernel's rounding
-modes and is not where the gain is. Two matmuls are bundled: `fp4_matmul` takes a
-numerically exact FP4 path (dequantize, bf16 accumulate) that is portable to every
-compute-capability-8.0+ GPU, and `fp4_mm` runs the operands on the Blackwell NVFP4
-tensor cores by swizzling the e4m3 block scales into the layout `torch._scaled_mm`
-consumes. On an RTX PRO 6000 the hardware path returns a bit-identical result
-(relative error 0) and is about 10-14x faster than the dequantize path and 1.5-1.7x
-faster than a bf16 matmul.
 ## Requirements and limits
-- NVIDIA GPU with compute capability 8.0+ (e4m3 scale via software conversion).
 - `K` a multiple of 16; float32 or bfloat16 input.
 ## References
 NVFP4 and OCP microscaling (MXFP) formats; Gupta et al., "Deep Learning with
-Limited Numerical Precision" (stochastic rounding, 2015); Seide et al., "1-bit
-SGD" (error feedback, 2014); low-precision training with unbiased quantization.
 ## License

 # fp4-train
+fp4-train is an NVFP4 quantizer with stochastic rounding and error feedback,
+the primitives for stable 4-bit training, loadable through `kernels`. The
+reference baseline is an independent NVFP4 encoding implementation, matched
+to zero error, with fp32 training as the convergence reference.
+Training with weights stored in 4-bit fails under ordinary rounding: once the
+per-step update is smaller than a quantization step, round-to-nearest
+discards every update and the model stops learning. This kernel quantizes to
+NVFP4 (e2m1 values with an e4m3 scale per 16-element block) with stochastic
+rounding, so quantization is unbiased and sub-step updates accumulate in
+expectation, and it carries a per-element error-feedback residual so the
+dropped fraction of every update is banked instead of lost. The packed output
+is the layout Blackwell FP4 tensor cores consume.
+![Four training curves: round-to-nearest flatlines while stochastic rounding and error feedback descend toward the fp32 reference, above an e2m1 grid where a value hops between neighbors](https://huggingface.co/kernels/phanerozoic/fp4-train/resolve/main/media/hero.gif)
+*Weights stored in NVFP4 and updated in place, 500 steps of sub-step updates:
+round-to-nearest never leaves its starting loss (6.56); stochastic rounding
+reaches 1.80, error feedback 0.77, against the fp32 reference at 0.02. Below,
+quantizing 2.3 on the e2m1 grid hops between 2 and 3 with the running mean at
+2.26.*
 ## Usage
 fp4 = get_kernel("phanerozoic/fp4-train", version=1, trust_remote_code=True)
+# NVFP4-stored-weight training: the weight lives in 4-bit, updated in place
 ef = fp4.ErrorFeedback(stochastic=False)                      # RTN + error feedback
+W = W0.clone()
 for step in range(num_steps):
+    grad = compute_grad(W)
+    W = ef.qdq("layer0.weight", W - lr * grad, step=step)
 # quantization primitives
+packed, scale = fp4.quantize(W, stochastic=True, seed=0, step=global_step)
 Wq = fp4.dequantize(packed, scale, *W.shape)
 # native NVFP4 tensor-core matmul on Blackwell (a @ b^T, K a multiple of 16)
+y = fp4.fp4_mm(a, b)
 ```
 `version` selects the release branch; `trust_remote_code` is required by
 | `quantize(x, stochastic, residual, seed, step)` | NVFP4 quantize; error feedback if `residual` is a `[M,K]` f32 tensor (updated in place) |
 | `dequantize(packed, scale, M, K)` | NVFP4 -> bf16 |
 | `qdq(x, stochastic, residual, seed, step)` | quantize then dequantize (the value the tensor cores use) |
+| `fp4_linear(x, weight, stochastic, seed, step)` | quantized-operand linear (straight-through gradient) |
 | `fp4_matmul(a, b, stochastic, seed, step)` | `a @ b^T` FP4-operand product, dequant + bf16 accumulate (portable, cc 8.0+) |
 | `fp4_mm(a, b, stochastic, seed, step)` | `a @ b^T` on the Blackwell NVFP4 tensor cores (native `_scaled_mm`, ~10x faster) |
 | `to_fp4_operands(x, stochastic, seed, step)` | packed fp4 + e4m3 scale as the torch FP4 dtypes for `_scaled_mm` |
 | `ErrorFeedback` | per-tensor residual state for weight quantization |
+## Method
+Each 16-element block gets a scale of `block_absmax / 6` encoded to e4m3, so
+the largest element maps near the e2m1 maximum. Every element is divided by
+the decoded scale and encoded to the e2m1 grid `{0, .5, 1, 1.5, 2, 3, 4, 6}`
+with a sign bit. Rounding is stochastic: for a value between two grid points,
+the higher is chosen with probability equal to the fractional distance, using
+a counter-based `(seed, step, global-index)` RNG, so the expected quantized
+value equals the input. With error feedback, the residual
+`x - dequant(quantize(x))` is carried and added into the next quantization,
+so the time-averaged value equals the true value even under round-to-nearest.
+Codes pack two per byte; block scales are e4m3, the layout NVFP4 tensor cores
+read.
 ## Correctness
 Measured on L4, H200, and RTX PRO 6000 through `get_kernel`:
+- Exact encoding: quantize/dequantize matches an independent reference to 0
+  (`max abs diff = 0`).
+- Stochastic rounding is unbiased: for 2.3 between grid points 2 and 3, the
+  SR mean is 2.300 where round-to-nearest is biased to 2.000.
+- Error feedback cancels the drift: RTN plus error feedback time-averages to
+  2.297.
+- With NVFP4-stored weights and sub-step updates, RTN stalls at the starting
+  loss while SR and EF recover near-full-precision loss: a linear model
+  reaches MSE 1.0 (EF) and 13 (SR) against bf16 0.0 with RTN at 77; a small
+  MLP reaches 0.159 (EF) and 0.171 (SR) against bf16 0.158 with RTN at 1.03.
+- On an RTX PRO 6000 the hardware `fp4_mm` path returns a bit-identical
+  result to the dequantize path (relative error 0) and is about 10-14x
+  faster, and 1.5-1.7x faster than a bf16 matmul.
 ## Requirements and limits
+- NVIDIA GPU with compute capability 8.0+ (e4m3 scale via software
+  conversion); `fp4_mm` requires Blackwell.
 - `K` a multiple of 16; float32 or bfloat16 input.
+- The rounding modes matter when the low precision is in the training state
+  itself (weights stored and updated in NVFP4). With a full-precision master
+  weight, the master carries the sub-ULP information and the choice of
+  rounding is second order.
+- Blocks whose absmax falls below the e4m3 scale range quantize to zero; keep
+  stored-weight scales inside the representable range.
 ## References
 NVFP4 and OCP microscaling (MXFP) formats; Gupta et al., "Deep Learning with
+Limited Numerical Precision" (stochastic rounding, 2015); Seide et al.,
+"1-bit SGD" (error feedback, 2014).
 ## License