fix(loader,eval): disable continuous-batching CUDA graphs for paged attention and require Evalution>=0.0.10 (#2978)

* fix(loader,eval): disable continuous-batching CUDA graphs for paged attention and require Evalution>=0.0.10

* chore(version): bump to 7.3.3

Author: Qubitium

gptqmodel/models/loader.py

@@ -178,6 +178,44 @@ def _override_attn_implementation(config: PretrainedConfig, attn_implementation:
             pass
 
 
+def _set_paged_attention_safe_cuda_graphs(model) -> None:
+    """Disable CUDA graph capture for paged attention by default.
+
+    ``transformers`` continuous batching defaults ``use_cuda_graph=True`` for flash-attention
+    paths, but ``flash_attn_with_kvcache`` is not currently CUDA-graph capture-safe. Setting
+    the model's default ``continuous_batching_config`` keeps paged Flash Attention working for
+    direct ``generate_batch`` / ``init_continuous_batching`` callers without requiring a monkey
+    patch.
+    """
+    if not hasattr(model, "config") or not isinstance(model.config, PretrainedConfig):
+        return
+
+    attn = getattr(model.config, "_attn_implementation", None) or getattr(
+        model.config, "attn_implementation", None
+    )
+    if not isinstance(attn, str) or "paged" not in [part.strip() for part in attn.split("|")]:
+        return
+
+    try:
+        from transformers import ContinuousBatchingConfig, GenerationConfig
+    except Exception:
+        return
+
+    gen_config = getattr(model, "generation_config", None)
+    if gen_config is None:
+        try:
+            gen_config = GenerationConfig.from_model_config(model.config)
+        except Exception:
+            return
+        model.generation_config = gen_config
+
+    cb_config = getattr(gen_config, "continuous_batching_config", None)
+    if not isinstance(cb_config, ContinuousBatchingConfig):
+        gen_config.continuous_batching_config = ContinuousBatchingConfig(use_cuda_graph=(False, False))
+    elif getattr(cb_config, "use_cuda_graph", None) is None:
+        cb_config.use_cuda_graph = (False, False)
+
+
 def _is_accelerated_attention_device(device: object) -> bool:
     """Return True when the selected device can run CUDA/ROCm flash attention."""
 
@@ -663,6 +701,7 @@ def from_pretrained(
                 trust_remote_code=trust_remote_code,
                 model_local_path=model_local_path,
             )
+            _set_paged_attention_safe_cuda_graphs(instance.model)
 
             return instance
 
@@ -825,6 +864,7 @@ def skip(*args, **kwargs):
             trust_remote_code=trust_remote_code,
             model_local_path=model_local_path,
         )
+        _set_paged_attention_safe_cuda_graphs(instance.model)
 
         timer = getattr(instance, "quant_region_timer", None)
         if timer is not None:
@@ -1078,6 +1118,7 @@ def from_quantized(
                 model_local_path=model_local_path,
             )
             instance._runtime_generate = runtime_generate
+            _set_paged_attention_safe_cuda_graphs(instance.model)
             return instance
 
         if format_code == FORMAT.MARLIN:
@@ -1677,7 +1718,7 @@ def assign(mod, device_id):
                 cls.generate = lambda _, **kwargs: mlx_generate(model=model, tokenizer=tokenizer, **kwargs)
 
 
-        return cls(
+        instance = cls(
             model,
             quantized=True,
             quantize_config=qcfg,
@@ -1687,6 +1728,8 @@ def assign(mod, device_id):
             trust_remote_code=trust_remote_code,
             model_local_path=model_local_path,
         )
+        _set_paged_attention_safe_cuda_graphs(instance.model)
+        return instance
 
     cls.from_quantized = from_quantized
 

gptqmodel/version.py

@@ -7,4 +7,4 @@
 # even minor versions are release
 # 5.2.0 => release, 5.1.0 => devel
 # micro version (5.2.x) denotes patch fix, i.e. 5.2.1 is a patch fix release
-__version__ = "7.3.2"
+__version__ = "7.3.3"

pyproject.toml

@@ -66,7 +66,7 @@ hf = [
     "optimum>=1.21.2",
 ]
 eval = [
-    "Evalution>=0.0.9",
+    "Evalution>=0.0.10",
 ]
 triton = [
     "triton>=3.4.0",

tests/eval.py

@@ -48,6 +48,7 @@
     "tokenizer_revision",
     "tp_size",
     "trust_remote_code",
+    "use_cuda_graph",
     "vllm_path",
 }
 _DROPPED_MODEL_ARG_KEYS = {
@@ -533,6 +534,7 @@ def _build_evalution_runtime(
             "padding_side": engine_padding_side,
             "backend": _normalize_backend_name(backend),
             "gptqmodel_path": str(Path(__file__).resolve().parents[2]),
+            "use_cuda_graph": engine_options.get("use_cuda_graph"),
         }
         engine = safe_kwargs_call(evalution.GPTQModel, kwargs=engine_kwargs)
         model_config = evalution.Model(