已合并
fix(triton_experimental): rebuild stale-rank store broadcast after r-tree promotion #44955
huyuchao创建于 17 天前
fix(triton_experimental): rebuild stale-rank store broadcast after r-tree promotion #44955
已合并
共 2 个文件变更+134-1
| @@ -0,0 +1,84 @@ | |||
| 1 | +# Owner(s): ["module: tests"] | ||
| 2 | +# Regression test for the stale-rank store broadcast on promoted r-trees. | ||
| 3 | +# | ||
| 4 | +# A whole-tensor reduction whose output is a single scalar (xnumel == 1) is | ||
| 5 | +# stored via an integer-index store. Upstream store codegen appends | ||
| 6 | +# ``.broadcast_to(<value.shape>)`` using the CSE shape captured BEFORE the | ||
| 7 | +# triton_experimental r-tree promotion raises the value's rank, so the store | ||
| 8 | +# line keeps its pre-promotion rank and Triton rejects the kernel with | ||
| 9 | +# ``ValueError('Cannot broadcast, rank mismatch')`` unless | ||
| 10 | +# _rewrite_reduction_store_shape rebuilds it. This exercises the CausalLM loss | ||
| 11 | +# pattern (shift + log_softmax + nll_loss with ignore_index) that produced | ||
| 12 | +# triton_unk_fused_clone_nll_loss_forward_slice_view_* on Electra/Roberta/XGLM. | ||
| 13 | + | ||
| 14 | +import torch | ||
| 15 | +import torch.nn.functional as F | ||
| 16 | +from torch._inductor.utils import run_and_get_code | ||
| 17 | +from torch.testing._internal.common_utils import ( | ||
| 18 | + run_tests, | ||
| 19 | + parametrize, | ||
| 20 | + instantiate_parametrized_tests, | ||
| 21 | +) | ||
| 22 | +from testutils import TestUtils | ||
| 23 | + | ||
| 24 | +import torch_npu # noqa: F401 | ||
| 25 | + | ||
| 26 | +# Heuristics import emitted only by the triton_experimental wrapper header | ||
| 27 | +# (torch_npu/_inductor/triton_experimental/codegen/wrapper.py): identifies | ||
| 28 | +# which codegen backend produced the wrapper. | ||
| 29 | +EXPERIMENTAL_MARKER = "triton_experimental import npu_triton_heuristics" | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +class TestPromotedRtreeScalarStore(TestUtils): | ||
| 33 | + | ||
| 34 | + def setUp(self): | ||
| 35 | + super().setUp() | ||
| 36 | + # Reset dynamo/inductor caches so the test forces fresh codegen. | ||
| 37 | + torch._dynamo.reset() | ||
| 38 | + | ||
| 39 | + def tearDown(self): | ||
| 40 | + torch._dynamo.reset() | ||
| 41 | + super().tearDown() | ||
| 42 | + | ||
| 43 | + | ||
| 44 | + def _causal_lm_loss(logits, labels): | ||
| 45 | + # Shift like a CausalLM: predict labels[:, 1:] from logits[:, :-1]. | ||
| 46 | + # labels[:, 1:] is a strided [B, S-1] slice and its reshape inserts a | ||
| 47 | + # clone into the graph, so the reduction index decomposes into two | ||
| 48 | + # r-nodes and the r-tree gets promoted; the loss is a single scalar | ||
| 49 | + # stored at a constant index. | ||
| 50 | + lsm = torch.log_softmax(logits[:, :-1, :].float(), dim=-1) | ||
| 51 | + return F.nll_loss( | ||
| 52 | + lsm.reshape(-1, lsm.shape[-1]), | ||
| 53 | + labels[:, 1:].reshape(-1), | ||
| 54 | + ignore_index=-100, | ||
| 55 | + ) | ||
| 56 | + | ||
| 57 | + | ||
| 58 | + def test_promoted_rtree_scalar_store(self, dtype): | ||
| 59 | + batch, seq, vocab = 32, 512, 32 | ||
| 60 | + logits = torch.randn( | ||
| 61 | + batch, seq, vocab, dtype=eval(f"torch.{dtype}"), device="npu" | ||
| 62 | + ) | ||
| 63 | + labels = torch.randint( | ||
| 64 | + 0, vocab, (batch, seq), dtype=torch.int64, device="npu" | ||
| 65 | + ) | ||
| 66 | + labels[:, :100] = -100 # exercise the ignore_index masking path | ||
| 67 | + | ||
| 68 | + eager_out = self._causal_lm_loss(logits, labels) | ||
| 69 | + | ||
| 70 | + compiled = torch.compile( | ||
| 71 | + self._causal_lm_loss, options={"npu_backend": "triton_experimental"} | ||
| 72 | + ) | ||
| 73 | + compiled_out, codes = run_and_get_code(compiled, logits, labels) | ||
| 74 | + | ||
| 75 | + # The compile must succeed (pre-fix it dies with the rank mismatch) | ||
| 76 | + # and produce the triton_experimental wrapper. | ||
| 77 | + self.assertIn(EXPERIMENTAL_MARKER, codes[0]) | ||
| 78 | + torch.testing.assert_close(eager_out, compiled_out, rtol=1e-4, atol=1e-4) | ||
| 79 | + | ||
| 80 | + | ||
| 81 | +instantiate_parametrized_tests(TestPromotedRtreeScalarStore) | ||
| 82 | + | ||
| 83 | +if __name__ == "__main__": | ||
| 84 | + run_tests() | ||
| @@ -543,7 +543,56 @@ def _npu_apply_promoted_rtree_lines( | |||
| 543 | None, | 543 | None, |
| 544 | ) | 544 | ) |
| 545 | if broadcast is None: | 545 | if broadcast is None: |
| 546 | - return raw | 546 | + # Stale pre-promotion rank: upstream store codegen appends |
| 547 | + # ``.broadcast_to(<value.shape>)`` using the CSE shape captured BEFORE | ||
| 548 | + # r-tree promotion, so the arg count can be smaller than real_ndim | ||
| 549 | + # (e.g. ``(XBLOCK, 1)`` against a rank-3 promoted value) and Triton | ||
| 550 | + # rejects the store ("Cannot broadcast, rank mismatch"). Rebuild the | ||
| 551 | + # args: r-slots forced to 1, the remaining slots take the old non-1 | ||
| 552 | + # args in order; also pad the base ``tl.full`` shape list so the | ||
| 553 | + # broadcast input rank matches (Triton requires equal ranks). | ||
| 554 | + cand = next( | ||
| 555 | + ( | ||
| 556 | + item | ||
| 557 | + for item in ast.walk(store) | ||
| 558 | + if isinstance(item, ast.Call) | ||
| 559 | + and isinstance(item.func, ast.Attribute) | ||
| 560 | + and item.func.attr == "broadcast_to" | ||
| 561 | + ), | ||
| 562 | + None, | ||
| 563 | + ) | ||
| 564 | + if cand is None or len(cand.args) >= real_ndim: | ||
| 565 | + return raw | ||
| 566 | + old_non1 = [ | ||
| 567 | + a | ||
| 568 | + for a in cand.args | ||
| 569 | + if not (isinstance(a, ast.Constant) and str(a.value) == "1") | ||
| 570 | + ] | ||
| 571 | + free_slots = [s for s in range(real_ndim) if s not in r_slots] | ||
| 572 | + if len(old_non1) != len(free_slots): | ||
| 573 | + return raw | ||
| 574 | + broadcast = cand | ||
| 575 | + new_args = [] | ||
| 576 | + for s in range(real_ndim): | ||
| 577 | + if s in r_slots: | ||
| 578 | + new_args.append( | ||
| 579 | + ast.copy_location(ast.Constant(value=1), broadcast) | ||
| 580 | + ) | ||
| 581 | + else: | ||
| 582 | + new_args.append(old_non1.pop(0)) | ||
| 583 | + broadcast.args[:] = new_args | ||
| 584 | + base = broadcast.func.value | ||
| 585 | + if ( | ||
| 586 | + isinstance(base, ast.Call) | ||
| 587 | + and base.args | ||
| 588 | + and isinstance(base.args[0], ast.List) | ||
| 589 | + and len(base.args[0].elts) < real_ndim | ||
| 590 | + ): | ||
| 591 | + elts = base.args[0].elts | ||
| 592 | + while len(elts) < real_ndim: | ||
| 593 | + elts.insert( | ||
| 594 | + 0, ast.copy_location(ast.Constant(value=1), base.args[0]) | ||
| 595 | + ) | ||
| 547 | for slot in r_slots: | 596 | for slot in r_slots: |
| 548 | if 0 <= slot < len(broadcast.args): | 597 | if 0 <= slot < len(broadcast.args): |
| 549 | broadcast.args[slot] = ast.copy_location( | 598 | broadcast.args[slot] = ast.copy_location( |