support import for recipe snippets

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Author: shengliangxu

modelopt/recipe/_config_loader.py

@@ -103,9 +103,39 @@ def load_config(config_file: str | Path | Traversable) -> dict[str, Any] | list[
             f"Cannot find config file of {config_file}, paths checked: {paths_to_check}"
         )
 
-    _raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
-    if _raw is None:
+    text = config_path.read_text(encoding="utf-8")
+    docs = list(yaml.safe_load_all(text))
+
+    if len(docs) == 0 or docs[0] is None:
         return {}
+    if len(docs) == 1:
+        _raw = docs[0]
+    elif len(docs) == 2:
+        # Multi-document: first doc is imports/metadata, second is content.
+        # Merge the imports into the content for downstream resolution.
+        header, content = docs[0], docs[1]
+        if not isinstance(header, dict):
+            raise ValueError(
+                f"Config file {config_path}: first YAML document must be a mapping, "
+                f"got {type(header).__name__}"
+            )
+        if content is None:
+            content = {}
+        if isinstance(content, dict):
+            _raw = {**header, **content}
+        elif isinstance(content, list):
+            # List content with a header dict — attach imports via wrapper
+            _raw = {**header, "_list_content": content}
+        else:
+            raise ValueError(
+                f"Config file {config_path}: second YAML document must be a mapping or list, "
+                f"got {type(content).__name__}"
+            )
+    else:
+        raise ValueError(
+            f"Config file {config_path}: expected 1 or 2 YAML documents, got {len(docs)}"
+        )
+
     if not isinstance(_raw, (dict, list)):
         raise ValueError(
             f"Config file {config_path} must contain a YAML mapping or list, got {type(_raw).__name__}"

modelopt/recipe/loader.py

@@ -79,6 +79,9 @@ def _resolve_imports(
         snippet = load_config(config_path)
         if isinstance(snippet, dict) and "imports" in snippet:
             snippet = _resolve_imports(snippet, _loading | {config_path})
+        # Unwrap _list_content (multi-document YAML: imports + list content)
+        if isinstance(snippet, dict) and "_list_content" in snippet:
+            snippet = snippet["_list_content"]
         import_map[name] = snippet
 
     def _lookup(ref_name: str, context: str) -> Any:
@@ -89,58 +92,65 @@ def _lookup(ref_name: str, context: str) -> Any:
             )
         return import_map[ref_name]
 
+    def _resolve_list(entries: list[Any]) -> list[Any]:
+        """Resolve $import markers in a list of quant_cfg-style entries."""
+        resolved: list[Any] = []
+        for entry in entries:
+            if isinstance(entry, dict) and _IMPORT_KEY in entry:
+                # {$import: name} → splice imported list
+                if len(entry) > 1:
+                    raise ValueError(
+                        f"$import must be the only key in the dict, got extra keys: "
+                        f"{sorted(k for k in entry if k != _IMPORT_KEY)}"
+                    )
+                imported = _lookup(entry[_IMPORT_KEY], "list entry")
+                if not isinstance(imported, list):
+                    raise ValueError(
+                        f"$import {entry[_IMPORT_KEY]!r} in list must resolve to a "
+                        f"list, got {type(imported).__name__}."
+                    )
+                resolved.extend(imported)
+            elif (
+                isinstance(entry, dict)
+                and isinstance(entry.get("cfg"), dict)
+                and _IMPORT_KEY in entry["cfg"]
+            ):
+                # cfg: {$import: name_or_list, ...inline} → import then override
+                #
+                # Precedence (lowest → highest):
+                #   1. Imports in list order (later imports override earlier)
+                #   2. Inline keys (override all imports)
+                ref = entry["cfg"].pop(_IMPORT_KEY)
+                inline_keys = dict(entry["cfg"])
+                ref_names = ref if isinstance(ref, list) else [ref]
+
+                merged: dict[str, Any] = {}
+                for name in ref_names:
+                    snippet = _lookup(name, f"cfg of {entry}")
+                    if not isinstance(snippet, dict):
+                        raise ValueError(
+                            f"$import {name!r} in cfg must resolve to a dict, "
+                            f"got {type(snippet).__name__}."
+                        )
+                    merged.update(snippet)
+
+                merged.update(inline_keys)
+                entry["cfg"] = merged
+                resolved.append(entry)
+            else:
+                resolved.append(entry)
+        return resolved
+
     # Resolve $import references in quant_cfg entries
     quantize = data.get("quantize")
     if isinstance(quantize, dict):
         quant_cfg = quantize.get("quant_cfg")
         if isinstance(quant_cfg, list):
-            resolved_cfg: list[Any] = []
-            for entry in quant_cfg:
-                if isinstance(entry, dict) and _IMPORT_KEY in entry:
-                    # {$import: name} → splice imported list into quant_cfg
-                    if len(entry) > 1:
-                        raise ValueError(
-                            f"$import must be the only key in the dict, got extra keys: "
-                            f"{sorted(k for k in entry if k != _IMPORT_KEY)}"
-                        )
-                    imported = _lookup(entry[_IMPORT_KEY], "quant_cfg entry")
-                    if not isinstance(imported, list):
-                        raise ValueError(
-                            f"$import {entry[_IMPORT_KEY]!r} in quant_cfg must resolve to a "
-                            f"list, got {type(imported).__name__}. Config snippets used as "
-                            f"quant_cfg entries must be YAML lists."
-                        )
-                    resolved_cfg.extend(imported)
-                elif (
-                    isinstance(entry, dict)
-                    and isinstance(entry.get("cfg"), dict)
-                    and _IMPORT_KEY in entry["cfg"]
-                ):
-                    # cfg: {$import: name_or_list, ...inline} → import then override
-                    #
-                    # Precedence (lowest → highest):
-                    #   1. Imports in list order (later imports override earlier)
-                    #   2. Inline keys (override all imports)
-                    ref = entry["cfg"].pop(_IMPORT_KEY)
-                    inline_keys = dict(entry["cfg"])  # remaining inline keys
-                    ref_names = ref if isinstance(ref, list) else [ref]
-
-                    merged: dict[str, Any] = {}
-                    for name in ref_names:
-                        snippet = _lookup(name, f"cfg of {entry}")
-                        if not isinstance(snippet, dict):
-                            raise ValueError(
-                                f"$import {name!r} in cfg must resolve to a dict, "
-                                f"got {type(snippet).__name__}."
-                            )
-                        merged.update(snippet)
-
-                    merged.update(inline_keys)
-                    entry["cfg"] = merged
-                    resolved_cfg.append(entry)
-                else:
-                    resolved_cfg.append(entry)
-            quantize["quant_cfg"] = resolved_cfg
+            quantize["quant_cfg"] = _resolve_list(quant_cfg)
+
+    # Resolve $import references in _list_content (multi-document snippets)
+    if "_list_content" in data:
+        data["_list_content"] = _resolve_list(data["_list_content"])
 
     return data
 

modelopt_recipes/configs/ptq/fp8_kv.yaml

@@ -0,0 +1,7 @@
+# FP8 E4M3 KV cache quantization.
+imports:
+  fp8: configs/numerics/fp8
+---
+  - quantizer_name: '*[kv]_bmm_quantizer'
+    cfg:
+      $import: fp8

modelopt_recipes/general/ptq/fp8_default-fp8_kv.yml

@@ -17,6 +17,7 @@ imports:
   base_disable_all: configs/ptq/base_disable_all
   default_disabled: configs/ptq/default_disabled_quantizers
   fp8: configs/numerics/fp8
+  fp8_kv: configs/ptq/fp8_kv
 
 metadata:
   recipe_type: ptq
@@ -31,7 +32,5 @@ quantize:
     - quantizer_name: '*weight_quantizer'
       cfg:
         $import: fp8
-    - quantizer_name: '*[kv]_bmm_quantizer'
-      cfg:
-        $import: fp8
+    - $import: fp8_kv
     - $import: default_disabled

modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yml

@@ -17,7 +17,7 @@ imports:
   base_disable_all: configs/ptq/base_disable_all
   default_disabled: configs/ptq/default_disabled_quantizers
   nvfp4: configs/numerics/nvfp4_dynamic
-  fp8: configs/numerics/fp8
+  fp8_kv: configs/ptq/fp8_kv
 
 metadata:
   recipe_type: ptq
@@ -32,7 +32,5 @@ quantize:
     - quantizer_name: '*input_quantizer'
       cfg:
         $import: nvfp4
-    - quantizer_name: '*[kv]_bmm_quantizer'
-      cfg:
-        $import: fp8
+    - $import: fp8_kv
     - $import: default_disabled

modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yml

@@ -17,7 +17,7 @@ imports:
   base_disable_all: configs/ptq/base_disable_all
   default_disabled: configs/ptq/default_disabled_quantizers
   nvfp4: configs/numerics/nvfp4_dynamic
-  fp8: configs/numerics/fp8
+  fp8_kv: configs/ptq/fp8_kv
 
 metadata:
   recipe_type: ptq
@@ -38,7 +38,5 @@ quantize:
     - quantizer_name: '*block_sparse_moe*input_quantizer'
       cfg:
         $import: nvfp4
-    - quantizer_name: '*[kv]_bmm_quantizer'
-      cfg:
-        $import: fp8
+    - $import: fp8_kv
     - $import: default_disabled

modelopt_recipes/general/ptq/nvfp4_mlp_only-fp8_kv.yml

@@ -17,7 +17,7 @@ imports:
   base_disable_all: configs/ptq/base_disable_all
   default_disabled: configs/ptq/default_disabled_quantizers
   nvfp4: configs/numerics/nvfp4_dynamic
-  fp8: configs/numerics/fp8
+  fp8_kv: configs/ptq/fp8_kv
 
 metadata:
   recipe_type: ptq
@@ -38,7 +38,5 @@ quantize:
     - quantizer_name: '*block_sparse_moe*input_quantizer'
       cfg:
         $import: nvfp4
-    - quantizer_name: '*[kv]_bmm_quantizer'
-      cfg:
-        $import: fp8
+    - $import: fp8_kv
     - $import: default_disabled

modelopt_recipes/general/ptq/nvfp4_omlp_only-fp8_kv.yml

@@ -17,7 +17,7 @@ imports:
   base_disable_all: configs/ptq/base_disable_all
   default_disabled: configs/ptq/default_disabled_quantizers
   nvfp4: configs/numerics/nvfp4_dynamic
-  fp8: configs/numerics/fp8
+  fp8_kv: configs/ptq/fp8_kv
 
 metadata:
   recipe_type: ptq
@@ -44,7 +44,5 @@ quantize:
     - quantizer_name: '*o_proj*input_quantizer'
       cfg:
         $import: nvfp4
-    - quantizer_name: '*[kv]_bmm_quantizer'
-      cfg:
-        $import: fp8
+    - $import: fp8_kv
     - $import: default_disabled

tests/unit/recipe/test_loader.py

@@ -650,6 +650,49 @@ def test_import_dir_format(tmp_path):
     assert recipe.quantize["quant_cfg"][0]["cfg"] == {"num_bits": (4, 3), "axis": None}
 
 
+# ---------------------------------------------------------------------------
+# imports — multi-document snippets
+# ---------------------------------------------------------------------------
+
+
+def test_import_multi_document_list_snippet(tmp_path):
+    """List snippet using multi-document YAML (imports --- content) resolves $import."""
+    (tmp_path / "fp8.yml").write_text("num_bits: e4m3\n")
+    (tmp_path / "kv.yaml").write_text(
+        f"imports:\n"
+        f"  fp8: {tmp_path / 'fp8.yml'}\n"
+        f"---\n"
+        f"- quantizer_name: '*[kv]_bmm_quantizer'\n"
+        f"  cfg:\n"
+        f"    $import: fp8\n"
+    )
+    recipe_file = tmp_path / "recipe.yml"
+    recipe_file.write_text(
+        f"imports:\n"
+        f"  kv: {tmp_path / 'kv.yaml'}\n"
+        f"metadata:\n"
+        f"  recipe_type: ptq\n"
+        f"quantize:\n"
+        f"  algorithm: max\n"
+        f"  quant_cfg:\n"
+        f"    - $import: kv\n"
+    )
+    recipe = load_recipe(recipe_file)
+    assert len(recipe.quantize["quant_cfg"]) == 1
+    assert recipe.quantize["quant_cfg"][0]["quantizer_name"] == "*[kv]_bmm_quantizer"
+    assert recipe.quantize["quant_cfg"][0]["cfg"] == {"num_bits": (4, 3)}
+
+
+def test_import_builtin_fp8_kv_snippet():
+    """Built-in fp8_kv snippet uses multi-document format and resolves correctly."""
+    recipe = load_recipe("general/ptq/fp8_default-fp8_kv")
+    kv_entries = [
+        e for e in recipe.quantize["quant_cfg"] if e.get("quantizer_name") == "*[kv]_bmm_quantizer"
+    ]
+    assert len(kv_entries) == 1
+    assert kv_entries[0]["cfg"]["num_bits"] == (4, 3)
+
+
 # ---------------------------------------------------------------------------
 # imports — recursive resolution and cycle detection
 # ---------------------------------------------------------------------------