[fix] multi graphs capture error

2025-09-18 16:00:49 +00:00 · 2024-03-11 10:49:31 +08:00
parent cefaeb5fdd
commit b2c0d9ff2b
4 changed files with 27 additions and 30 deletions
--- a/colossalai/kernel/triton/rms_layernorm.py
+++ b/colossalai/kernel/triton/rms_layernorm.py
@@ -92,7 +92,6 @@ if HAS_TRITON:

    def rms_layernorm(x, weight, eps, norm_output=None, residual=None):
        # allocate output
-        # y = torch.empty_like(x) if norm_output is None else norm_output
        y = (
            x * 0 if norm_output is None else norm_output
        )  # to make the operation non-functional, store y as the intermediate activation