From 4c98759c73a9b2a657290149b5f30e438f0ae7eb Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Fri, 21 Aug 2026 16:13:21 +0800 Subject: [PATCH] [AMD] fix(rocm): support flydsl 0.3.0 in the FlyDSL fused norm kernel (#34536) Co-authored-by: Bingxu Chen <195740905+bingxche@users.noreply.github.com> Co-authored-by: thomawan --- .../ops/diffusion/norm/fused_residual_norm_flydsl.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py index f7c666d14..cf5e45428 100644 --- a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py +++ b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py @@ -26,10 +26,17 @@ from flydsl._mlir.dialects import ( ) from flydsl._mlir.dialects import vector as _vector from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr +from flydsl.expr import arith, const_expr, range_constexpr from flydsl.expr.arith import ArithValue, CmpIPredicate from flydsl.expr.typing import Int32, T +try: + from flydsl.expr import buffer_ops +except ImportError: + # flydsl 0.3.0 removed flydsl.expr.buffer_ops and pushed the buffer-resource + # layer down to its consumers; AITER carries the same module, same API. + from aiter.ops.flydsl.kernels import buffer_ops + WARP_SIZE = 64 _VEC = 8 _NUM_WAVES = 10