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

[AMD] Fix CI multimodal-gen-test-1-gpu-amd for gen model (#21621)

Browse Source
This commit is contained in:
jacky.cheng
2026-03-30 23:02:20 -07:00
committed by GitHub
parent 03e4f2858d
commit 8ba992411d
1 changed files with 28 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
+28 -14
View File
@@ -79,13 +79,20 @@ def _fused_layernorm_scale_shift_gate_select01_kernel(
shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs) # Branch on scalar idx instead of using tl.where on pointers.
shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs) # tl.where on pointers triggers an assertion in AMD Triton's
gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs) # CanonicalizePointers pass (ConvertArithSelectOp) on gfx950.
# This keeps it at 3 loads (not 6), avoids the pointer-level
scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32) # tl.where entirely, and since idx is uniform across all threads
shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32) # the branch has no divergence cost.
gate = tl.load(gate_ptrs, mask=mask, other=0.0) if idx:
scale = tl.load(scale1_ptrs, mask=mask, other=0.0).to(tl.float32)
shift = tl.load(shift1_ptrs, mask=mask, other=0.0).to(tl.float32)
gate = tl.load(gate1_ptrs, mask=mask, other=0.0)
else:
scale = tl.load(scale0_ptrs, mask=mask, other=0.0).to(tl.float32)
shift = tl.load(shift0_ptrs, mask=mask, other=0.0).to(tl.float32)
gate = tl.load(gate0_ptrs, mask=mask, other=0.0)
y = x_hat * (1.0 + scale) + shift y = x_hat * (1.0 + scale) + shift
tl.store(out_row_ptr + cols, y, mask=mask) tl.store(out_row_ptr + cols, y, mask=mask)
@@ -180,13 +187,20 @@ def _fused_residual_layernorm_scale_shift_gate_select01_kernel(
shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs) # Branch on scalar idx instead of using tl.where on pointers.
shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs) # tl.where on pointers triggers an assertion in AMD Triton's
gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs) # CanonicalizePointers pass (ConvertArithSelectOp) on gfx950.
# This keeps it at 3 loads (not 6), avoids the pointer-level
scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32) # tl.where entirely, and since idx is uniform across all threads
shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32) # the branch has no divergence cost.
gate = tl.load(gate_ptrs, mask=mask, other=0.0) if idx:
scale = tl.load(scale1_ptrs, mask=mask, other=0.0).to(tl.float32)
shift = tl.load(shift1_ptrs, mask=mask, other=0.0).to(tl.float32)
gate = tl.load(gate1_ptrs, mask=mask, other=0.0)
else:
scale = tl.load(scale0_ptrs, mask=mask, other=0.0).to(tl.float32)
shift = tl.load(shift0_ptrs, mask=mask, other=0.0).to(tl.float32)
gate = tl.load(gate0_ptrs, mask=mask, other=0.0)
y = x_hat * (1.0 + scale) + shift y = x_hat * (1.0 + scale) + shift
tl.store(out_row_ptr + cols, y, mask=mask) tl.store(out_row_ptr + cols, y, mask=mask)