From 5e4baea64d0f4cd1b6ce147e4f7a28f0f41d995f Mon Sep 17 00:00:00 2001 From: Lei Ke <1141466880@qq.com> Date: Wed, 9 Sep 2026 00:08:46 +0800 Subject: [PATCH] [Bugfix][NPU] Limit MiniMax H3 modulation grid size (#6794) Signed-off-by: KrystalRay Co-authored-by: KrystalRay --- .buildkite/npu/test-npu-ready.yml | 1 + .../minimax_h3/test_minimax_h3_modulation.py | 142 +++++++++++++++++- .../attention/ops/minimax_h3_modulation.py | 48 +++++- 3 files changed, 182 insertions(+), 9 deletions(-) diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 1f25ebac1..01fb32879 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -20,6 +20,7 @@ steps: tests/diffusion/layers/test_activation.py tests/diffusion/layers/test_fused_qk_norm_rope_npu.py tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py + tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py -m 'core_model and npu and A3' --run-level 'core_model' diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index dbeeaa66a..e10fbe56a 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -5,14 +5,154 @@ import pytest import torch from vllm.triton_utils import HAS_TRITON +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( + _MAX_1D_GRID_SIZE, + _iter_row_chunks, + _launch_row_chunks, + indexed_gate, indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, rms_norm_indexed_scale_shift, ) +from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +@pytest.mark.cpu +@pytest.mark.parametrize( + ("rows", "expected"), + [ + (0, []), + (_MAX_1D_GRID_SIZE, [(0, _MAX_1D_GRID_SIZE)]), + (_MAX_1D_GRID_SIZE + 1, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 1)]), + (69888, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 4353)]), + ], +) +def test_row_chunks_respect_ascend_grid_limit(rows, expected): + assert list(_iter_row_chunks(rows)) == expected + + +@pytest.mark.cpu +@pytest.mark.parametrize( + ("device_type", "expected_grids_and_offsets"), + [ + ("npu", [(_MAX_1D_GRID_SIZE, 0), (4353, _MAX_1D_GRID_SIZE)]), + ("cuda", [(69888, 0)]), + ], +) +def test_launch_row_chunks_is_npu_only(device_type, expected_grids_and_offsets): + class RecordingKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.calls.append((grid, args, kwargs)) + + return launch + + kernel = RecordingKernel() + first_arg = object() + second_arg = object() + + _launch_row_chunks(kernel, 69888, device_type, first_arg, second_arg, block_n=16) + + assert kernel.calls == [ + ((grid,), (first_arg, second_arg, row_offset), {"block_n": 16}) + for grid, row_offset in expected_grids_and_offsets + ] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + "device", + [pytest.param("npu", marks=hardware_marks(res={"npu": "A3"}, num_cards=1))], +) +def test_modulation_wrappers_match_reference_above_grid_limit(device): + rows = 69888 + hidden_size = 16 + num_conditions = 7 + eps = 1e-6 + dtype = torch.bfloat16 + + torch.manual_seed(0) + x = torch.randn(rows, hidden_size, dtype=dtype) + residual = torch.randn_like(x) + branch = torch.randn_like(x) + weight = torch.randn(hidden_size, dtype=dtype) + shift = torch.randn(num_conditions, hidden_size, dtype=dtype) + scale = torch.randn_like(shift) + gate = torch.randn_like(shift) + indices = torch.arange(rows) % num_conditions + + indexed_shift = shift.index_select(0, indices) + indexed_scale = scale.index_select(0, indices) + indexed_gate_values = gate.index_select(0, indices) + + x_float = x.float() + branch_float = branch.float() + indexed_shift_float = indexed_shift.float() + indexed_scale_float = indexed_scale.float() + indexed_gate_float = indexed_gate_values.float() + + expected_scale_shift = (x_float * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + expected_gate = (x_float + indexed_gate_float * branch_float).to(dtype) + + normalized = x_float + normalized = normalized * torch.rsqrt(normalized.pow(2).mean(-1, keepdim=True) + eps) + normalized = normalized * weight.float() + expected_rms = (normalized * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + + updated_residual = residual.float() + indexed_gate_float * branch_float + expected_residual = updated_residual.to(dtype) + normalized_residual = updated_residual + normalized_residual = normalized_residual * torch.rsqrt(normalized_residual.pow(2).mean(-1, keepdim=True) + eps) + normalized_residual = normalized_residual * weight.float() + expected_modulated = (normalized_residual * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + + npu_x = x.to(device) + npu_residual = residual.to(device) + npu_branch = branch.to(device) + npu_weight = weight.to(device) + npu_shift = shift.to(device) + npu_scale = scale.to(device) + npu_gate = gate.to(device) + npu_indices = indices.to(device) + + actual_scale_shift = indexed_scale_shift_(npu_x.clone(), npu_shift, npu_scale, npu_indices) + actual_gate = indexed_gate(npu_x, npu_gate, npu_branch, npu_indices) + actual_rms = rms_norm_indexed_scale_shift( + npu_x, + npu_weight, + npu_shift, + npu_scale, + npu_indices, + eps, + ) + actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift( + npu_residual, + npu_gate, + npu_branch, + npu_weight, + npu_shift, + npu_scale, + npu_indices, + eps, + ) + + for actual, expected in ( + (actual_scale_shift, expected_scale_shift), + (actual_gate, expected_gate), + (actual_rms, expected_rms), + (actual_residual, expected_residual), + (actual_modulated, expected_modulated), + ): + torch.testing.assert_close(actual.cpu(), expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.cuda @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") def test_fused_modulation_preserves_bf16_residual_boundary() -> None: diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index f75833831..20327b7e9 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -7,6 +7,22 @@ from __future__ import annotations import torch from vllm.triton_utils import tl, triton +# Ascend CANN limits the launch coreDim to uint16. +_MAX_1D_GRID_SIZE = 65535 + + +def _iter_row_chunks(rows: int): + for row_offset in range(0, rows, _MAX_1D_GRID_SIZE): + yield row_offset, min(rows - row_offset, _MAX_1D_GRID_SIZE) + + +def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) -> None: + if device_type != "npu": + kernel[(rows,)](*args, 0, **kwargs) + return + for row_offset, chunk_rows in _iter_row_chunks(rows): + kernel[(chunk_rows,)](*args, row_offset, **kwargs) + @triton.jit def _indexed_scale_shift_kernel( @@ -20,9 +36,10 @@ def _indexed_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -51,9 +68,10 @@ def _indexed_gate_kernel( stride_gate_row, stride_other_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -83,9 +101,10 @@ def _rms_norm_indexed_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -123,9 +142,10 @@ def _indexed_gate_rms_norm_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -162,7 +182,10 @@ def indexed_scale_shift_( rows, hidden_size = x.shape if rows == 0: return x - _indexed_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _indexed_scale_shift_kernel, + rows, + x.device.type, x, x, shift, @@ -192,7 +215,10 @@ def indexed_gate( rows, hidden_size = x.shape if rows == 0: return output - _indexed_gate_kernel[(rows,)]( + _launch_row_chunks( + _indexed_gate_kernel, + rows, + x.device.type, output, x, gate, @@ -229,7 +255,10 @@ def rms_norm_indexed_scale_shift( output = torch.empty_like(x) rows, hidden_size = x.shape if rows: - _rms_norm_indexed_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _rms_norm_indexed_scale_shift_kernel, + rows, + x.device.type, output, x, weight, @@ -274,7 +303,10 @@ def indexed_gate_rms_norm_scale_shift( modulated_out = torch.empty_like(residual) rows, hidden_size = residual.shape if rows: - _indexed_gate_rms_norm_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _indexed_gate_rms_norm_scale_shift_kernel, + rows, + residual.device.type, residual_out, modulated_out, residual,