[Bugfix][NPU] Limit MiniMax H3 modulation grid size (#6794)

Signed-off-by: KrystalRay <keeleiray@gmail.com>
Co-authored-by: KrystalRay <keeleiray@gmail.com>

.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'
 

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:

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,