diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py index f964c1577..4e433b108 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py @@ -11,26 +11,18 @@ if HAS_TRITON: @triton.jit def _qk_rms_norm_rope_exact_kernel( - q_ptr, - k_ptr, + x_ptr, cos_ptr, sin_ptr, - q_output_ptr, - k_output_ptr, - q_stride_token, - q_stride_head, - q_stride_dim, - k_stride_token, - k_stride_head, - k_stride_dim, + output_ptr, + x_stride_token, + x_stride_head, + x_stride_dim, rope_stride_token, rope_stride_dim, - q_output_stride_token, - q_output_stride_head, - q_output_stride_dim, - k_output_stride_token, - k_output_stride_head, - k_output_stride_dim, + output_stride_token, + output_stride_head, + output_stride_dim, num_heads: tl.constexpr, head_dim: tl.constexpr, rotary_dim: tl.constexpr, @@ -39,15 +31,12 @@ if HAS_TRITON: ): token = tl.program_id(0) head_group = tl.program_id(1) - qk_domain = tl.program_id(2) heads = head_group * heads_per_program + tl.arange(0, heads_per_program) dims = tl.arange(0, head_dim) mask = heads[:, None] < num_heads - q_offsets = token * q_stride_token + heads[:, None] * q_stride_head + dims[None, :] * q_stride_dim - k_offsets = token * k_stride_token + heads[:, None] * k_stride_head + dims[None, :] * k_stride_dim - x_ptrs = tl.where(qk_domain == 0, q_ptr + q_offsets, k_ptr + k_offsets) - x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + offsets = token * x_stride_token + heads[:, None] * x_stride_head + dims[None, :] * x_stride_dim + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) normalized = (x * inv_rms[:, None]).to(tl.float16) @@ -57,10 +46,8 @@ if HAS_TRITON: dims + rotary_half, tl.where(dims < rotary_dim, dims - rotary_half, dims), ) - q_pair_offsets = token * q_stride_token + heads[:, None] * q_stride_head + pair_dims[None, :] * q_stride_dim - k_pair_offsets = token * k_stride_token + heads[:, None] * k_stride_head + pair_dims[None, :] * k_stride_dim - pair_x_ptrs = tl.where(qk_domain == 0, q_ptr + q_pair_offsets, k_ptr + k_pair_offsets) - pair_x = tl.load(pair_x_ptrs, mask=mask, other=0.0).to(tl.float32) + pair_offsets = token * x_stride_token + heads[:, None] * x_stride_head + pair_dims[None, :] * x_stride_dim + pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) pair_normalized = (pair_x * inv_rms[:, None]).to(tl.float16) rope_mask = dims < rotary_dim @@ -92,18 +79,10 @@ if HAS_TRITON: ) output = tl.where(rope_mask[None, :], rotated, normalized) - q_output_offsets = ( - token * q_output_stride_token + heads[:, None] * q_output_stride_head + dims[None, :] * q_output_stride_dim + output_offsets = ( + token * output_stride_token + heads[:, None] * output_stride_head + dims[None, :] * output_stride_dim ) - k_output_offsets = ( - token * k_output_stride_token + heads[:, None] * k_output_stride_head + dims[None, :] * k_output_stride_dim - ) - output_ptrs = tl.where( - qk_domain == 0, - q_output_ptr + q_output_offsets, - k_output_ptr + k_output_offsets, - ) - tl.store(output_ptrs, output, mask=mask) + tl.store(output_ptr + output_offsets, output, mask=mask) def _supported_inputs( @@ -164,11 +143,15 @@ def try_qk_norm_rope_exact( sin = sin.reshape(tokens, sin.shape[-1]) q_output = torch.empty_like(q) k_output = torch.empty_like(k) - # The first two launch dimensions are part of the exactness contract on - # the validated SM90 and SM103 targets: changing the warp-to-row mapping - # changes the FP32 RMS reduction order. + # This launch layout is part of the exactness contract on the validated + # SM90 and SM103 targets: changing the warp-to-row mapping changes the + # FP32 RMS reduction order. + # Keep Q and K in separate launches so each kernel has fixed input and + # output pointers. A combined qk-domain launch introduces runtime pointer + # selection; on SM90, its generated arithmetic differs enough to + # accumulate through the 50-step I2VA denoising loop. heads_per_program = 8 - grid = (tokens, triton.cdiv(heads, heads_per_program), 2) + grid = (tokens, triton.cdiv(heads, heads_per_program)) launch_args = { "num_heads": heads, "head_dim": head_dim, @@ -179,22 +162,29 @@ def try_qk_norm_rope_exact( } _qk_rms_norm_rope_exact_kernel[grid]( q, - k, cos, sin, q_output, - k_output, q.stride(0), q.stride(1), q.stride(2), - k.stride(0), - k.stride(1), - k.stride(2), cos.stride(0), cos.stride(1), q_output.stride(0), q_output.stride(1), q_output.stride(2), + **launch_args, + ) + _qk_rms_norm_rope_exact_kernel[grid]( + k, + cos, + sin, + k_output, + k.stride(0), + k.stride(1), + k.stride(2), + cos.stride(0), + cos.stride(1), k_output.stride(0), k_output.stride(1), k_output.stride(2),