Skip to content

fix(engines/base): default use_cuda_graph=False for paged attention continuous batching - #129

Merged
Qubitium merged 2 commits into
mainfrom
devin/fix-paged-cuda-graph-default
Jul 25, 2026
Merged

fix(engines/base): default use_cuda_graph=False for paged attention continuous batching#129
Qubitium merged 2 commits into
mainfrom
devin/fix-paged-cuda-graph-default

Conversation

@Qubitium

@Qubitium Qubitium commented Jul 25, 2026

Copy link
Copy Markdown
Contributor

Summary

transformers continuous batching defaults use_cuda_graph=True for flash-attention paths. When attn_implementation contains the paged marker (e.g. "paged|flash_attention_2"), ModelRunner captures a CUDA graph around the model forward, and the paged decode path calls flash_paged._paged_decode_forwardflash_attn_with_kvcacheflash_attn_gpu.fwd_kvcache, which is not capture-safe. It fails with CUDA error: operation not permitted when stream is capturing.

This change adds a __post_init__ to BaseEnginePagedBatchingConfig so that when use_cuda_graph is not explicitly set and attn_implementation contains a paged token, it defaults to (False, False). Paged Flash Attention still runs, but CUDA graph capture is disabled.

def __post_init__(self) -> None:
    if self.use_cuda_graph is not None:
        return
    attn = getattr(self, "attn_implementation", None)
    if isinstance(attn, str) and "paged" in [part.strip() for part in attn.split("|")]:
        self.use_cuda_graph = (False, False)

This keeps the fix in one place (Evalution) so GPTQModel/Transformers engine configs always get the safe default, and downstream callers don't need to set it themselves.

Version bump

  • pyproject.toml: 0.0.90.0.10
  • evalution/version.py fallback: 0.0.90.0.10

Verification

  • ruff check evalution/engines/base.py passes.
  • PYTHON_GIL=0 pytest -q tests/test_transformer.py passes (62/62).
  • PYTHON_GIL=0 pytest -q tests/test_gptqmodel_engine.py -k "not generate_and_score_on_cuda" passes; the skipped integration test fails because gptqmodel is not importable under PYTHON_GIL=0 in this environment, which is unrelated to this change.
  • End-to-end gsm8k_platinum_cot (64 samples) on a pre-quantized Llama 3.2 1B GPTQ checkpoint with attn_implementation=paged|flash_attention_2 runs backend=continuous_batching paged_attention=True and completes with no CUDA graph-capture error.

@Qubitium Qubitium self-assigned this Jul 25, 2026
@devin-ai-integration

Copy link
Copy Markdown

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment, CI, and merge conflict monitoring

@Qubitium
Qubitium merged commit 19192b4 into main Jul 25, 2026
2 checks passed
@Qubitium
Qubitium deleted the devin/fix-paged-cuda-graph-default branch July 25, 2026 04:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant