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
`_power_of_two_ceiling`	Round positive normal FP32 values up to an exactly represented power of two.
`quantize_cache`	Apply the released model’s quantize/dequantize cache representation.

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._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.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