MCPcopy
hub / github.com/z-lab/dflash / _get_transformers_attn_impl

Function _get_transformers_attn_impl

dflash/benchmark.py:185–195  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

183
184
185def _get_transformers_attn_impl() -> str:
186 try:
187 import flash_attn # noqa: F401
188 return "flash_attention_2"
189 except ImportError:
190 logger.warning(
191 "flash_attn not installed. Falling back to torch.sdpa. Speedup will be lower. "
192 "For optimal speedup in Transformers backend, please install: "
193 "pip install flash-attn --no-build-isolation"
194 )
195 return "sdpa"
196
197
198def _run_transformers(args: argparse.Namespace) -> None:

Callers 1

_run_transformersFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected