Parse a torchao.dequantize_nvfp4 node. Returns (qdata, scale, per_tensor_scale, output_dtype) or None if not a dequantize_nvfp4 node or the custom op is not registered.
(
node: Node,
)
| 282 | |
| 283 | |
| 284 | def parse_dequant_nvfp4_node( |
| 285 | node: Node, |
| 286 | ) -> Optional[Tuple[Node, Node, Node, torch.dtype]]: |
| 287 | """Parse a torchao.dequantize_nvfp4 node. |
| 288 | |
| 289 | Returns (qdata, scale, per_tensor_scale, output_dtype) or None if not a |
| 290 | dequantize_nvfp4 node or the custom op is not registered. |
| 291 | """ |
| 292 | target = get_aten_target(node.target) |
| 293 | try: |
| 294 | import executorch.extension.llm.export.nvfp4 # noqa: F401 |
| 295 | except ImportError: |
| 296 | return None |
| 297 | |
| 298 | if target is not torch.ops.torchao.dequantize_nvfp4.default: |
| 299 | return None |
| 300 | |
| 301 | qdata, scale, per_tensor_scale = node.args[0:3] |
| 302 | |
| 303 | output_dtype = torch.float32 |
| 304 | if len(node.args) > 4: |
| 305 | output_dtype = node.args[4] |
| 306 | elif "output_dtype" in node.kwargs: |
| 307 | output_dtype = node.kwargs["output_dtype"] |
| 308 | |
| 309 | return qdata, scale, per_tensor_scale, output_dtype |
| 310 | |
| 311 | |
| 312 | def parse_dequant_node( |
no test coverage detected