vllm.model_executor.layers.fusion.relu2_nvfp4_quant

Functions:

• relu_squared_nvfp4_quant –

  Fuse BF16 ReLU2 and block16 NVFP4 with 128x4 swizzled scales.

relu_squared_nvfp4_quant(x, linear)

Fuse BF16 ReLU2 and block16 NVFP4 with 128x4 swizzled scales.

Source code in vllm/model_executor/layers/fusion/relu2_nvfp4_quant.py

def relu_squared_nvfp4_quant(
    x: torch.Tensor, linear: LinearBase
) -> QuantizedActivation:
    """Fuse BF16 ReLU2 and block16 NVFP4 with 128x4 swizzled scales."""
    from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
    from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic

    output, block_scale = torch.ops.vllm.relu_squared_nvfp4_quant(
        x, linear.input_global_scale_inv
    )
    return QuantizedActivation(
        data=output,
        scale=block_scale,
        orig_dtype=x.dtype,
        orig_shape=x.shape,
        quant_key=kNvfp4Dynamic,
    )