feat: Named graph specializations in specializations.json (Prefill/Decode/Vision/Encoder/Embedding) (quic#904)

## Summary
The backend compiler team requested a new specializations.json format
where each entry carries a meaningful graph name (e.g. "Prefill",
"Decode")

  ## Changes
- **`QEfficient/utils/_utils.py`** — new `_infer_specialization_name()`
and `to_named_specializations()` helpers
- **`QEfficient/base/modeling_qeff.py`** — `_compile()` uses new format
  - **`QEfficient/compile/qnn_compiler.py`** — QNN path uses new format
- **`QEfficient/compile/compile_helper.py`** — legacy
`create_and_dump_specializations()` uses new format

  ## Name inference rules
  | Keys present | Assigned name |
  |---|---|
  | `vision_size` / `img_size` / `grid_*`, no `seq_len` | `Vision` |
  | `encoder_ctx_len`, no `seq_len` | `Encoder` |
  | `sequence_length`, no `seq_len` | `Embedding` |
  | `seq_len != 1` | `Prefill` |
  | `seq_len == 1` | `Decode` |
  | anything else | `Graph_N` |

## Testing
21-unit tests added to `tests/unit_test/models/test_model_quickcheck.py`
covering causal LM, continuous batching, VLM vision/language, Whisper,
encoder/decoder, text embedding, and end-to-end JSON roundtrip.

cc: @anujgupt-github @quic-rishinr

---------

Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
Signed-off-by: Rishin Raj <rishinr@qti.qualcomm.com>
Co-authored-by: Rishin Raj <rishinr@qti.qualcomm.com>

Author: vbaddi

QEfficient/base/modeling_qeff.py

@@ -41,6 +41,7 @@
     hash_dict_params,
     load_json,
     require_value,
+    to_named_specializations,
 )
 from QEfficient.utils.export_utils import export_wrapper
 
@@ -471,6 +472,7 @@ def _compile(
         enable_chunking: Optional[bool] = False,
         retain_full_kv: Optional[bool] = None,
         qaic_config: Optional[dict] = None,
+        specialization_module_name: Optional[str] = None,
         **compiler_options,
     ) -> str:
         """
@@ -609,7 +611,7 @@ def _compile(
         if specializations is not None:
             specializations_json = compile_dir / "specializations.json"
             specializations_data = {
-                "specializations": [{k: str(v) for k, v in spec.items()} for spec in specializations]
+                "specializations": to_named_specializations(specializations, module_name=specialization_module_name)
             }
             create_json(str(specializations_json), specializations_data)
             command.append(f"-network-specialization-config={specializations_json}")

QEfficient/compile/compile_helper.py

@@ -14,35 +14,34 @@
 
 from QEfficient.compile.qnn_compiler import compile as qnn_compile
 from QEfficient.utils import constants
-from QEfficient.utils._utils import load_json, load_yaml
+from QEfficient.utils._utils import load_json, load_yaml, to_named_specializations
 from QEfficient.utils.logging_utils import logger
 
 
 def create_and_dump_specializations(
     batch_size: int, prompt_len: int, ctx_len: int, path: str, full_batch_size: Optional[int] = None
 ):
-    # Create specialization file.
-    specializations = {
-        "specializations": [
-            {
-                "batch_size": str(batch_size),
-                "seq_len": str(prompt_len),
-                "ctx_len": str(ctx_len),
-            },
-            {"batch_size": str(batch_size), "seq_len": "1", "ctx_len": str(ctx_len)},
-        ]
-    }
-    # If continuous batching is enabled by proving full_batch_size we need to add FBS to the specialization file and update the batch size of decoder part to FBS
+    # Build the base specialization entries first, then convert to named format.
+    base_specializations = [
+        {
+            "batch_size": str(batch_size),
+            "seq_len": str(prompt_len),
+            "ctx_len": str(ctx_len),
+        },
+        {"batch_size": str(batch_size), "seq_len": "1", "ctx_len": str(ctx_len)},
+    ]
+    # If continuous batching is enabled by providing full_batch_size we need to add FBS to the specialization file and update the batch size of decoder part to FBS
     if full_batch_size is not None:
-        specializations["specializations"][0]["full_batch_size"] = str(full_batch_size)
-        specializations["specializations"][1]["full_batch_size"] = str(full_batch_size)
-        specializations["specializations"][1]["batch_size"] = str(full_batch_size)
+        base_specializations[0]["full_batch_size"] = str(full_batch_size)
+        base_specializations[1]["full_batch_size"] = str(full_batch_size)
+        base_specializations[1]["batch_size"] = str(full_batch_size)
 
-    # To handle repetative input in specializations when prompt_len is 1
+    # To handle repetitive input in specializations when prompt_len is 1
     if prompt_len == 1 and full_batch_size is None:
-        specializations["specializations"].pop()
+        base_specializations.pop()
 
     # Dump
+    specializations = {"specializations": to_named_specializations(base_specializations)}
     with open(path, "w") as file:
         json.dump(specializations, file, indent=4)
 

QEfficient/compile/qnn_compiler.py

@@ -11,7 +11,7 @@
 import shutil
 from typing import Dict, List, Optional
 
-from QEfficient.utils._utils import create_json, execute_command, load_json
+from QEfficient.utils._utils import create_json, execute_command, load_json, to_named_specializations
 from QEfficient.utils.constants import QnnConstants
 from QEfficient.utils.generate_qnn_network_specialization_config import (
     generate_data_format_config,
@@ -423,7 +423,7 @@ def compile(
             specializations_json = qpc_base_path / "specializations.json"
             with open(specializations_json, "w") as fp:
                 json.dump(
-                    {"specializations": [{k: str(v) for k, v in spec.items()} for spec in specializations]},
+                    {"specializations": to_named_specializations(specializations)},
                     fp,
                     indent=4,
                 )

QEfficient/diffusers/pipelines/pipeline_utils.py

@@ -173,6 +173,12 @@ def _prepare_and_compile(module_name: str, module_obj: Any) -> None:
         else:
             specializations = [specializations]
 
+        # Tag each spec with the module name so _compile knows the graph name.
+        specializations = [
+            {**s, "_graph_name": f"{module_name}_model_type_{s['model_type']}" if "model_type" in s else module_name}
+            for s in specializations
+        ]
+
         if module_obj.qpc_path is None:
             # Compile with prepared specializations
             module_obj.compile(specializations=specializations, **compile_kwargs)
@@ -226,6 +232,12 @@ def compile_modules_sequential(
         else:
             specializations = [specializations]
 
+        # Tag each spec with the module name so _compile knows the graph name.
+        specializations = [
+            {**s, "_graph_name": f"{module_name}_model_type_{s['model_type']}" if "model_type" in s else module_name}
+            for s in specializations
+        ]
+
         if module_obj.qpc_path is None:
             # Compile with prepared specializations
             module_obj.compile(specializations=specializations, **compile_kwargs)

QEfficient/generation/text_generation_inference.py

@@ -188,9 +188,12 @@ def get_compilation_dims(qpc_path: str) -> Tuple[int, int, Optional[int]]:
     else:
         raise FileNotFoundError(f"expected specializations.json file at path, {qpc_base_path}")
 
-    compilation_batch_size = int(data["specializations"][0]["batch_size"])
-    compilation_ctx_len = int(data["specializations"][0]["ctx_len"])
-    if compilation_fbs := data["specializations"][0].get("full_batch_size", None):
+    # Support both the legacy flat format and the new {name, symbols} format.
+    first = data["specializations"][0]
+    spec = first.get("symbols", first)
+    compilation_batch_size = int(spec["batch_size"])
+    compilation_ctx_len = int(spec["ctx_len"])
+    if compilation_fbs := spec.get("full_batch_size", None):
         compilation_fbs = int(compilation_fbs)
     return compilation_batch_size, compilation_ctx_len, compilation_fbs
 

QEfficient/transformers/models/modeling_auto.py

@@ -449,8 +449,14 @@ def compile(
         if isinstance(seq_len, list) and len(seq_len) >= 15:
             warnings.warn("Recommended: `seq_len` should contain fewer than 15 items.")
 
+        _seq_lens = seq_len if isinstance(seq_len, list) else [seq_len]
         specializations = [
-            {"batch_size": batch_size, "seq_len": sl} for sl in (seq_len if isinstance(seq_len, list) else [seq_len])
+            {
+                "_graph_name": "Embedding" if len(_seq_lens) == 1 else f"Embedding_{i}",
+                "batch_size": batch_size,
+                "seq_len": sl,
+            }
+            for i, sl in enumerate(_seq_lens)
         ]
 
         target_dtype = getattr(self.model.config, "torch_dtype", torch.float32)
@@ -794,8 +800,14 @@ def compile(
         if isinstance(seq_len, list) and len(seq_len) >= 15:
             warnings.warn("Recommended: `seq_len` should contain fewer than 15 items.")
 
+        _seq_lens = seq_len if isinstance(seq_len, list) else [seq_len]
         specializations = [
-            {"batch_size": batch_size, "seq_len": sl} for sl in (seq_len if isinstance(seq_len, list) else [seq_len])
+            {
+                "_graph_name": "SeqClassification" if len(_seq_lens) == 1 else f"SeqClassification_{i}",
+                "batch_size": batch_size,
+                "seq_len": sl,
+            }
+            for i, sl in enumerate(_seq_lens)
         ]
         target_dtype = getattr(self.model.config, "torch_dtype", torch.float32)
         return self._compile(
@@ -1582,6 +1594,7 @@ def compile(
                 compile_dir=compile_dir,
                 compile_only=True,
                 specializations=specializations["vision"],
+                specialization_module_name="Vision",
                 convert_to_fp16=(CUSTOM_IO_DTYPE_MAP[target_dtype] == "float16"),
                 mxfp6_matmul=constants.VISION_MXFP6_MATMUL,
                 mdp_ts_num_devices=num_devices,
@@ -3256,7 +3269,9 @@ def build_prefill_specialization(
         # TODO: remove this; not required
         if full_batch_size:
             spec["full_batch_exec_size"] = exec_batch_size
-        return {k: v for k, v in spec.items() if v is not None}
+        result = {k: v for k, v in spec.items() if v is not None}
+        result["_graph_name"] = "Prefill"
+        return result
 
     def build_decode_specialization(
         self,
@@ -3314,7 +3329,9 @@ def build_decode_specialization(
             spec["full_batch_size"] = kv_cache_batch_size
         else:
             spec["batch_size"] = kv_cache_batch_size
-        return {k: v for k, v in spec.items() if v is not None}
+        result = {k: v for k, v in spec.items() if v is not None}
+        result["_graph_name"] = "Decode"
+        return result
 
     def compile(
         self,
@@ -4235,8 +4252,10 @@ def compile(
             :str: Path of the compiled ``qpc`` package.
         """
 
+        _seq_lens = seq_len if isinstance(seq_len, list) else [seq_len]
         specializations = [
-            {"batch_size": batch_size, "seq_len": sl} for sl in (seq_len if isinstance(seq_len, list) else [seq_len])
+            {"_graph_name": "CTC" if len(_seq_lens) == 1 else f"CTC_{i}", "batch_size": batch_size, "seq_len": sl}
+            for i, sl in enumerate(_seq_lens)
         ]
 
         target_dtype = getattr(self.model.config, "torch_dtype", torch.float32)

QEfficient/transformers/models/whisper/modeling_whisper.py

@@ -835,6 +835,7 @@ def get_specializations(self, batch_size: int, encoder_ctx_len, ctx_len, **compi
         feature_len = encoder_ctx_len * 2
 
         encoder_specializations = {
+            "_graph_name": "Encoder",
             "batch_size": batch_size,
             "seq_len": 1,
             "encoder_ctx_len": encoder_ctx_len,
@@ -843,6 +844,7 @@ def get_specializations(self, batch_size: int, encoder_ctx_len, ctx_len, **compi
         }
 
         decoder_specializations = {
+            "_graph_name": "Decode",
             "batch_size": batch_size,
             "seq_len": 1,
             "encoder_ctx_len": encoder_ctx_len,