|
| 1 | +"""Opt-in switches for the Fun-ASR-Nano LLM forward, read from ``llm_conf``; both default to off. |
| 2 | +
|
| 3 | +``llm_conf.sdpa_backends``: a list of ``flash`` / ``efficient`` / ``math``. The decoder forward runs |
| 4 | +under ``torch.nn.attention.sdpa_kernel`` with these backends: the process-wide SDPA flags are set on |
| 5 | +entry and restored on return, so another thread calling scaled_dot_product_attention meanwhile sees |
| 6 | +the same selection (leave it unset when several models share one process concurrently). |
| 7 | +``llm_conf.torch_compile``: run the decoder stack through ``torch.compile(dynamic=True)`` for |
| 8 | +inputs on a CUDA device, decided per call; a CPU decoder and a 1-sequence batch stay eager. |
| 9 | +""" |
| 10 | + |
| 11 | +import torch |
| 12 | + |
| 13 | +SDPA_BACKENDS = {"flash": "FLASH_ATTENTION", "efficient": "EFFICIENT_ATTENTION", "math": "MATH"} |
| 14 | + |
| 15 | + |
| 16 | +def resolve_sdpa_backends(names): |
| 17 | + """``llm_conf.sdpa_backends`` (a list of names, or ``None``) -> ``SDPBackend`` members.""" |
| 18 | + if names is None: |
| 19 | + return None |
| 20 | + from torch.nn.attention import SDPBackend # a pybind enum: getattr, not subscript |
| 21 | + |
| 22 | + backends = [] |
| 23 | + for name in names: |
| 24 | + member = SDPA_BACKENDS.get(name.lower()) |
| 25 | + if member is None: |
| 26 | + raise ValueError( |
| 27 | + f"llm_conf.sdpa_backends: unknown backend {name!r}, " |
| 28 | + f"choose from {sorted(SDPA_BACKENDS)}" |
| 29 | + ) |
| 30 | + backends.append(getattr(SDPBackend, member)) |
| 31 | + return backends |
| 32 | + |
| 33 | + |
| 34 | +def configure_llm_forward(llm, llm_conf): |
| 35 | + """Install the switches ``llm_conf`` asks for on ``llm.model.forward`` (nothing by default).""" |
| 36 | + backends = resolve_sdpa_backends(llm_conf.get("sdpa_backends", None)) |
| 37 | + compile_requested = bool(llm_conf.get("torch_compile", False)) |
| 38 | + if backends is None and not compile_requested: |
| 39 | + return |
| 40 | + |
| 41 | + decoder = llm.model |
| 42 | + forward = eager = decoder.forward |
| 43 | + |
| 44 | + if compile_requested: |
| 45 | + compiled = torch.compile(eager, dynamic=True) |
| 46 | + |
| 47 | + def forward(*args, **kw): |
| 48 | + # Per call, not at construction: the trainer builds the model on the CPU and moves it |
| 49 | + # to the GPU afterwards. A 1-sequence batch would get its own specialised graph. |
| 50 | + x = kw.get("inputs_embeds", kw.get("input_ids")) |
| 51 | + if x is None and args: |
| 52 | + x = args[0] |
| 53 | + if x is None or x.device.type != "cuda" or x.shape[0] == 1: |
| 54 | + return eager(*args, **kw) |
| 55 | + return compiled(*args, **kw) |
| 56 | + |
| 57 | + if backends is not None: |
| 58 | + from torch.nn.attention import sdpa_kernel |
| 59 | + |
| 60 | + inner = forward |
| 61 | + |
| 62 | + def forward(*args, **kw): |
| 63 | + with sdpa_kernel(backends): # wraps the compiled call, so Dynamo traces under it |
| 64 | + return inner(*args, **kw) |
| 65 | + |
| 66 | + decoder.forward = forward |
0 commit comments