nemo_automodel.components.models.deepseek_v41.quantization

Reference cache QAT formats with straight-through activation gradients.

Module Contents

Classes

Name	Description
`_CacheQuantization`	Quantize cache values in forward and pass the activation gradient unchanged.

Functions

Name	Description
`_cache_quantization_core`	Round independent channel groups of `values`; the math behind `quantize_cache`.
`_power_of_two_ceiling`	Round positive normal FP32 values up to an exactly represented power of two.
`compile_cache_quantization`	Wrap the cache-quantization core with `torch.compile` (`DeepseekV41TextConfig.compile_cache_quant`).
`quantize_cache`	Apply the released model’s quantize/dequantize cache representation.

Data

_QUANT_COMPILED

_fp4_table_cache

_quantize

API

class nemo_automodel.components.models.deepseek_v41.quantization._CacheQuantization()

Bases: Function

Quantize cache values in forward and pass the activation gradient unchanged.

nemo_automodel.components.models.deepseek_v41.quantization._CacheQuantization.backward(
    ctx: object,
    gradient: torch.Tensor
) -> tuple[torch.Tensor, None, None]

staticmethod

Pass the gradient for […, channels] activations through quantization.

Parameters:

ctx

object

Unused forward context.

gradient

torch.Tensor

Tensor of shape […, channels] from the cache consumer.

Returns: torch.Tensor

Gradient of shape […, channels], followed by no gradients for the

nemo_automodel.components.models.deepseek_v41.quantization._CacheQuantization.forward(
    ctx: object,
    values: torch.Tensor,
    block_size: int,
    format: str
) -> torch.Tensor

staticmethod

Round independent channel groups without modifying the input.

Parameters:

ctx

object

Autograd context; no tensors need saving for the straight-through gradient.

values

torch.Tensor

Tensor of shape […, channels], with arbitrary leading dimensions.

block_size

int

Consecutive channels sharing a quantization scale.

format

str

fp8, mxfp4 (E8M0 scale), or nvfp4 (E4M3 scale).

Returns: torch.Tensor

Independently stored tensor of shape […, channels] with input dtype.

nemo_automodel.components.models.deepseek_v41.quantization._cache_quantization_core(
    values: torch.Tensor,
    block_size: int,
    format: str
) -> torch.Tensor

Round independent channel groups of values; the math behind quantize_cache.

Kept as a plain function so compile_cache_quantization can wrap it with torch.compile: eager, the fp8 path is five elementwise kernels and the FP4 paths about a dozen (pad, cast, amax, scale, clamp, bucketize, gather, copysign, two casts) per cache tensor.

Parameters:

values

torch.Tensor

Tensor of shape […, channels] with channels a multiple of block_size.

block_size

int

Consecutive channels sharing a quantization scale.

format

str

fp8, mxfp4 (E8M0 scale), or nvfp4 (E4M3 scale).

Returns: torch.Tensor

Independently stored tensor of shape […, channels] with the input dtype.

nemo_automodel.components.models.deepseek_v41.quantization._power_of_two_ceiling(
    value: torch.Tensor
) -> torch.Tensor

Round positive normal FP32 values up to an exactly represented power of two.

Parameters:

value

torch.Tensor

Positive FP32 tensor of arbitrary shape.

Returns: torch.Tensor

FP32 tensor of the same shape with integer powers of two, matching the

nemo_automodel.components.models.deepseek_v41.quantization.compile_cache_quantization() -> None

Wrap the cache-quantization core with torch.compile (DeepseekV41TextConfig.compile_cache_quant).

Runs once per process; static shapes (dynamic=False): one specialization per call site shape and format (SWA KV fp8, index Q/K mxfp4, compressed KV nvfp4). No fullgraph: past dynamo’s recompile limit a shape runs eager instead of raising. Numerics: the same operations fused, so results match eager up to summation order (there is none: every step is elementwise or a per-group max).

nemo_automodel.components.models.deepseek_v41.quantization.quantize_cache(
    values: torch.Tensor,
    format: typing.Literal['fp8', 'mxfp4', 'nvfp4'],
    block_size: int
) -> torch.Tensor

Apply the released model’s quantize/dequantize cache representation.

Parameters:

values

torch.Tensor

Tensor of shape […, channels], with arbitrary leading dimensions. Complete groups are required by the released kernels; a final partial group is zero-padded internally for scaled unit-test configurations.

format

Literal['fp8', 'mxfp4', 'nvfp4']

fp8 for SWA KV, mxfp4 for index Q/K, nvfp4 for compressed KV.

block_size

int

Channels per scale: 32 for SWA/indexer, 16 for compressed KV.

Returns: torch.Tensor

Tensor of shape […, channels], in the original dtype, with a

nemo_automodel.components.models.deepseek_v41.quantization._QUANT_COMPILED = False

nemo_automodel.components.models.deepseek_v41.quantization._fp4_table_cache: dict[tuple[device, dtype], tuple[Tensor, Tensor]] = {}

nemo_automodel.components.models.deepseek_v41.quantization._quantize = _cache_quantization_core