已合并
triton_experimental: softmax aclnn routing + npu_expand recursion fix #44889
huyuchao创建于 8月19日
triton_experimental: softmax aclnn routing + npu_expand recursion fix #44889
已合并
共 4 个文件变更+186-1
| @@ -364,6 +364,115 @@ def _override_softmax_backward_decomp_no_fma(): | |||
| 364 | _ind_decomps[aten._softmax_backward_data.default] = _softmax_backward_data_no_fma | 364 | _ind_decomps[aten._softmax_backward_data.default] = _softmax_backward_data_no_fma |
| 365 | 365 | ||
| 366 | 366 | ||
| 367 | +def _override_safe_softmax_decomp(): | ||
| 368 | + """Rewrite the mask-composite softmax (``aten._safe_softmax``, produced by | ||
| 369 | + the SDPA-with-mask decomposition) down to plain ``aten._softmax``. | ||
| 370 | + | ||
| 371 | + ``_safe_softmax``'s only semantic delta is the all-``-inf``-row guard | ||
| 372 | + (``where(all(row == -inf), zeros, softmax)``); in this backend's input | ||
| 373 | + domain that guard is dead code — attention masks never fully mask a row | ||
| 374 | + (causal keeps the diagonal; transformers rewrites fully-masked rows via | ||
| 375 | + ``_unmask_unattended`` before the mask reaches SDPA, see pytorch#110213). | ||
| 376 | + Dropping it removes the eager IsNegInf→All→SWhere chain (~5 ms/iter on | ||
| 377 | + TrOCR) and hands the op to the width router | ||
| 378 | + (``_override_plain_softmax_width_decomp``: Triton persistent fusion for | ||
| 379 | + narrow rows, aclnnSoftmax for wide rows). | ||
| 380 | + | ||
| 381 | + The optional ``dtype`` follows upstream ``torch.softmax(self, dim, dtype)`` | ||
| 382 | + semantics: cast before compute.""" | ||
| 383 | + from .triton_experimental import config as te_cfg | ||
| 384 | + | ||
| 385 | + if not te_cfg.safe_softmax_aclnn_fallback: | ||
| 386 | + return | ||
| 387 | + | ||
| 388 | + from torch._inductor.lowering import lowerings | ||
| 389 | + | ||
| 390 | + def _safe_softmax_to_softmax(x, dim, dtype=None): | ||
| 391 | + if dtype is not None: | ||
| 392 | + x = x.to(dtype) | ||
| 393 | + return aten._softmax(x, dim, False) | ||
| 394 | + | ||
| 395 | + decompositions[aten._safe_softmax.default] = _safe_softmax_to_softmax | ||
| 396 | + decompositions.pop(aten._safe_softmax, None) | ||
| 397 | + lowerings.pop(aten._safe_softmax, None) | ||
| 398 | + lowerings.pop(aten._safe_softmax.default, None) | ||
| 399 | + | ||
| 400 | + | ||
| 401 | +def _override_plain_softmax_width_decomp(): | ||
| 402 | + """Route plain ``aten._softmax`` by reduction width at decomposition time. | ||
| 403 | + | ||
| 404 | + Upstream lowers softmax purely through its decomposition (amax/exp/sum | ||
| 405 | + fused by the scheduler into a persistent-reduction kernel). Rows wider | ||
| 406 | + than the persistent-reduction budget lose that eligibility while TE keeps | ||
| 407 | + ``split_reductions`` off, degrading to a serial per-row scan — 2.5-2.8x | ||
| 408 | + slower than aclnnSoftmax at vocab widths (30k/50k, msprof on 910B2). | ||
| 409 | + ``_softmax.default`` is a leaf op (direct NPU kernel = aclnnSoftmax, no | ||
| 410 | + C++ composite key), so for wide rows emit it as a RAW node — temporarily | ||
| 411 | + lifting this wrapper out of the decomposition tables the tracer consults — | ||
| 412 | + and the node survives tracing; the fallback lowering installed below then | ||
| 413 | + routes it to aclnnSoftmax at runtime. Narrow rows keep the upstream | ||
| 414 | + decomposition verbatim (fusion with neighbours is profitable, | ||
| 415 | + BertForMaskedLM). Dynamic/symbolic widths fall back to the upstream | ||
| 416 | + decomposition (conservative).""" | ||
| 417 | + from .triton_experimental import config as te_cfg | ||
| 418 | + | ||
| 419 | + orig = decompositions.get(aten._softmax.default) | ||
| 420 | + if orig is None: | ||
| 421 | + return | ||
| 422 | + | ||
| 423 | + def _softmax_width_routed(x, dim, half_to_float=False): | ||
| 424 | + bound = te_cfg.softmax_aclnn_max_fuse_numel | ||
| 425 | + if bound > 0: | ||
| 426 | + try: | ||
| 427 | + sizes = x.size() | ||
| 428 | + width = sizes[dim if dim >= 0 else len(sizes) + dim] | ||
| 429 | + except Exception: | ||
| 430 | + width = None | ||
| 431 | + if isinstance(width, int) and width > bound: | ||
| 432 | + # Emit a raw _softmax.default node: without a table entry the | ||
| 433 | + # tracer keeps the leaf op intact (decompose() likewise | ||
| 434 | + # returns NotImplemented for it). Lift the wrapper out of | ||
| 435 | + # every table the tracer consults — the proxy mode's own | ||
| 436 | + # decomposition_table (the snapshot AOT built the mode with), | ||
| 437 | + # plus the live dict and select_decomp_table() — and restore | ||
| 438 | + # them once the node is traced. | ||
| 439 | + from torch._inductor.decomposition import select_decomp_table | ||
| 440 | + from torch.fx.experimental.proxy_tensor import get_proxy_mode | ||
| 441 | + tables = {id(decompositions): decompositions, | ||
| 442 | + id(t := select_decomp_table()): t} | ||
| 443 | + try: | ||
| 444 | + mode = get_proxy_mode() | ||
| 445 | + except Exception: | ||
| 446 | + mode = None | ||
| 447 | + if mode is not None: | ||
| 448 | + mt = getattr(mode, "decomposition_table", None) | ||
| 449 | + if mt is not None: | ||
| 450 | + tables[id(mt)] = mt | ||
| 451 | + saved = {} | ||
| 452 | + try: | ||
| 453 | + for tbl in tables.values(): | ||
| 454 | + saved[id(tbl)] = tbl.pop(aten._softmax.default, None) | ||
| 455 | + return aten._softmax(x, dim, half_to_float) | ||
| 456 | + finally: | ||
| 457 | + for tbl in tables.values(): | ||
| 458 | + s = saved.get(id(tbl)) | ||
| 459 | + if s is not None: | ||
| 460 | + tbl[aten._softmax.default] = s | ||
| 461 | + return orig(x, dim, half_to_float) | ||
| 462 | + | ||
| 463 | + decompositions[aten._softmax.default] = _softmax_width_routed | ||
| 464 | + | ||
| 465 | + # Wide-row raw nodes reach lowering with no entry of their own; unknown-op | ||
| 466 | + # handling calls make_fallback, which asserts against an op having BOTH a | ||
| 467 | + # fallback and a decomposition (we keep the decomposition for narrow | ||
| 468 | + # rows). Install the fallback lowering directly — narrow rows are always | ||
| 469 | + # decomposed at trace time and never consume it. | ||
| 470 | + from torch._inductor.lowering import lowerings, fallback_handler | ||
| 471 | + lowerings[aten._softmax.default] = fallback_handler( | ||
| 472 | + aten._softmax.default, add_to_fallback_set=True | ||
| 473 | + ) | ||
| 474 | + | ||
| 475 | + | ||
| 367 | def _override_gelu_decomp(): | 476 | def _override_gelu_decomp(): |
| 368 | from torch._inductor.decomposition import decompositions as _ind_decomps | 477 | from torch._inductor.decomposition import decompositions as _ind_decomps |
| 369 | 478 | ||
| @@ -507,6 +616,8 @@ def _register_triton_experimental_decompositions(): | |||
| 507 | 616 | ||
| 508 | _override_matmul_should_fold_for_npu() | 617 | _override_matmul_should_fold_for_npu() |
| 509 | _override_softmax_backward_decomp_no_fma() | 618 | _override_softmax_backward_decomp_no_fma() |
| 619 | + _override_safe_softmax_decomp() | ||
| 620 | + _override_plain_softmax_width_decomp() | ||
| 510 | _override_gelu_decomp() | 621 | _override_gelu_decomp() |
| 511 | _override_rms_norm_decomp() | 622 | _override_rms_norm_decomp() |
| 512 | _override_native_dropout_decomp() | 623 | _override_native_dropout_decomp() |
| @@ -85,8 +85,12 @@ def _activate(): | |||
| 85 | 85 | ||
| 86 | device.register_device_op_overrides_for_npu() | 86 | device.register_device_op_overrides_for_npu() |
| 87 | 87 | ||
| 88 | - from .lowering import _register_npu_inductor_fallbacks | 88 | + from .lowering import ( |
| 89 | + _register_npu_inductor_fallbacks, | ||
| 90 | + _register_softmax_aclnn_fallback, | ||
| 91 | + ) | ||
| 89 | _register_npu_inductor_fallbacks() | 92 | _register_npu_inductor_fallbacks() |
| 93 | + _register_softmax_aclnn_fallback() | ||
| 90 | 94 | ||
| 91 | device.register_interface_for_npu() | 95 | device.register_interface_for_npu() |
| 92 | 96 | ||
| @@ -172,6 +172,28 @@ refactor_clamp_stride: bool = False | |||
| 172 | # no-realize path. | 172 | # no-realize path. |
| 173 | realize_permute_gather: bool = True | 173 | realize_permute_gather: bool = True |
| 174 | 174 | ||
| 175 | +# Route the MASK-COMPOSITE softmax (aten._safe_softmax, produced from | ||
| 176 | +# transformers-style causal-mask + softmax patterns) through aclnn instead of | ||
| 177 | +# Triton fusion: the fused Triton kernel materializes [B,H,S,S] masks and | ||
| 178 | +# measured ~2x slower than the eager aclnn sequence (TrOCR: 10.3 ms/iter of | ||
| 179 | +# mask-materialization device time). Eager dispatch of _safe_softmax runs the | ||
| 180 | +# composite down to plain softmax → aclnnSoftmax (bit-identical, verified). | ||
| 181 | +# Plain softmax KEEPS its Triton fusion (profitable for e.g. BertForMaskedLM). | ||
| 182 | +safe_softmax_aclnn_fallback: bool = True | ||
| 183 | + | ||
| 184 | +# Route the log_softmax family through aclnnLogSoftmax: the loss-path Triton | ||
| 185 | +# log_softmax kernel is slower than aclnnLogSoftmax and its fusion drags a | ||
| 186 | +# gather decomposition along (TrOCR: Gather_AsStrided +2.9 ms/iter). | ||
| 187 | +log_softmax_aclnn_fallback: bool = True | ||
| 188 | + | ||
| 189 | +# Plain-softmax size routing: rows (reduction width) up to this bound keep the | ||
| 190 | +# Triton fusion; wider rows fall back to aclnnSoftmax (see decomposition.py | ||
| 191 | +# _override_plain_softmax_width_decomp: >1024 loses persistent-reduction | ||
| 192 | +# eligibility and TE keeps split_reductions off, wide rows degrade to a serial | ||
| 193 | +# per-row scan, 2.5-2.8x slower at vocab widths). 0 disables routing (always | ||
| 194 | +# Triton). | ||
| 195 | +softmax_aclnn_max_fuse_numel: int = 256 | ||
| 196 | + | ||
| 175 | # Reduction-tree real-block promotion (nested scalar r-loops -> real-block tile). | 197 | # Reduction-tree real-block promotion (nested scalar r-loops -> real-block tile). |
| 176 | rtree_real_block: bool = True | 198 | rtree_real_block: bool = True |
| 177 | 199 | ||
| @@ -48,6 +48,33 @@ npu = torch.ops.npu | |||
| 48 | FALLBACK_LIST = [] | 48 | FALLBACK_LIST = [] |
| 49 | 49 | ||
| 50 | 50 | ||
| 51 | +def _register_softmax_aclnn_fallback(): | ||
| 52 | + """See config.log_softmax_aclnn_fallback. | ||
| 53 | + | ||
| 54 | + Routed to aclnn (aten fallback): | ||
| 55 | + * ``aten.log_softmax`` family — the loss-path log_softmax Triton kernel is | ||
| 56 | + slower than aclnnLogSoftmax and its fusion drags a gather decomposition | ||
| 57 | + along (TrOCR: Gather_AsStrided +2.9 ms/iter). | ||
| 58 | + | ||
| 59 | + ``aten._safe_softmax`` (mask-composite) and the plain-softmax width | ||
| 60 | + routing are handled at loader time in ``torch_npu/_inductor/ | ||
| 61 | + decomposition.py`` — see ``_override_safe_softmax_decomp`` and | ||
| 62 | + ``_override_plain_softmax_width_decomp``.""" | ||
| 63 | + def _fallback(op_packet): | ||
| 64 | + for ov in ( | ||
| 65 | + (op_packet, op_packet.default) | ||
| 66 | + if hasattr(op_packet, "default") | ||
| 67 | + else (op_packet,) | ||
| 68 | + ): | ||
| 69 | + decompositions.pop(ov, None) | ||
| 70 | + lowerings.pop(ov, None) | ||
| 71 | + make_fallback(ov) | ||
| 72 | + | ||
| 73 | + if ncfg.log_softmax_aclnn_fallback: | ||
| 74 | + _fallback(aten.log_softmax) | ||
| 75 | + _fallback(aten._log_softmax) | ||
| 76 | + | ||
| 77 | + | ||
| 51 | def _register_npu_inductor_fallbacks(): | 78 | def _register_npu_inductor_fallbacks(): |
| 52 | gen_set = set() | 79 | gen_set = set() |
| 53 | for fn in GENERATE_LIST: | 80 | for fn in GENERATE_LIST: |
| @@ -357,6 +384,19 @@ def _is_tail_axis_broadcast(x, sizes): | |||
| 357 | 384 | ||
| 358 | from torch._inductor.lowering import expand as _orig_expand | 385 | from torch._inductor.lowering import expand as _orig_expand |
| 359 | 386 | ||
| 387 | +# Stash the pristine upstream handler on the (never-reloaded) upstream module: | ||
| 388 | +# re-executing THIS module would otherwise bind ``_orig_expand`` to the already | ||
| 389 | +# patched ``npu_expand`` (set below via ``lowering.expand = npu_expand``) and | ||
| 390 | +# recurse into itself forever (RecursionError seen in accuracy flows). | ||
| 391 | +_lowering_mod = torch._inductor.lowering | ||
| 392 | +if getattr(_lowering_mod, "_npu_upstream_expand_saved", None) is None: | ||
| 393 | + _lowering_mod._npu_upstream_expand_saved = _orig_expand | ||
| 394 | +_orig_expand = _lowering_mod._npu_upstream_expand_saved | ||
| 395 | + | ||
| 396 | +import threading as _threading | ||
| 397 | + | ||
| 398 | +_expand_reentry = _threading.local() | ||
| 399 | + | ||
| 360 | 400 | ||
| 361 | def npu_expand(x, sizes): | 401 | def npu_expand(x, sizes): |
| 362 | """NPU expand: realize a SHORT tail-axis broadcast into its own contiguous | 402 | """NPU expand: realize a SHORT tail-axis broadcast into its own contiguous |
| @@ -367,9 +407,15 @@ def npu_expand(x, sizes): | |||
| 367 | return result | 407 | return result |
| 368 | if not _is_tail_axis_broadcast(x, sizes): | 408 | if not _is_tail_axis_broadcast(x, sizes): |
| 369 | return result | 409 | return result |
| 410 | + if getattr(_expand_reentry, "active", False): | ||
| 411 | + # Re-entered while a tail-bcast realization is still materializing its | ||
| 412 | + # inputs (lazy values can drag lowering back through broadcast chains). | ||
| 413 | + # Keep the plain upstream result so the cycle terminates. | ||
| 414 | + return result | ||
| 370 | # realize() on the ExpandView realizes its underlying storage, not the [s,c] | 415 | # realize() on the ExpandView realizes its underlying storage, not the [s,c] |
| 371 | # result; wrap it in a Pointwise copy and realize THAT to materialize the | 416 | # result; wrap it in a Pointwise copy and realize THAT to materialize the |
| 372 | # broadcast as its own buffer. | 417 | # broadcast as its own buffer. |
| 418 | + _expand_reentry.active = True | ||
| 373 | try: | 419 | try: |
| 374 | copied = Pointwise.create( | 420 | copied = Pointwise.create( |
| 375 | device=result.get_device(), | 421 | device=result.get_device(), |
| @@ -381,6 +427,8 @@ def npu_expand(x, sizes): | |||
| 381 | except Exception as e: | 427 | except Exception as e: |
| 382 | log.debug("[NPU] realize tail-bcast expand skipped: %r", e) # noqa: G200 | 428 | log.debug("[NPU] realize tail-bcast expand skipped: %r", e) # noqa: G200 |
| 383 | return result | 429 | return result |
| 430 | + finally: | ||
| 431 | + _expand_reentry.active = False | ||
| 384 | if name: | 432 | if name: |
| 385 | V.graph.no_fuse_buffer_names.add(name) | 433 | V.graph.no_fuse_buffer_names.add(name) |
| 386 | log.debug("[NPU] realize tail-axis broadcast expand: buf=%s in_size=%s target=%s", | 434 | log.debug("[NPU] realize tail-axis broadcast expand: buf=%s in_size=%s target=%s", |