minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] fix(rocm): support flydsl 0.3.0 in the FlyDSL fused norm kernel (#34536)

Browse Source 
Co-authored-by: Bingxu Chen <195740905+bingxche@users.noreply.github.com>
Co-authored-by: thomawan <thomawan@amd.com>
This commit is contained in:
Bingxu Chen
2026-08-21 01:13:21 -07:00
committed by GitHub
co-authored by Bingxu Chen thomawan
parent 8ff9c2b227
commit 4c98759c73
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py
+8 -1
View File
@@ -26,10 +26,17 @@ from flydsl._mlir.dialects import (
) )
from flydsl._mlir.dialects import vector as _vector from flydsl._mlir.dialects import vector as _vector
from flydsl.compiler.kernel_function import CompilationContext 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.arith import ArithValue, CmpIPredicate
from flydsl.expr.typing import Int32, T 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 WARP_SIZE = 64
_VEC = 8 _VEC = 8
_NUM_WAVES = 10 _NUM_WAVES = 10