@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 68363d64ec

[ROCm][Perf] Avoid extra reshape kernel in Qwen GDN output norm (#47842)

Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Mikko Tukiainen committed Sep 22, 2026 at 09:38 UTC 68363d64ec476bbbf5025b30ced0f6a2425b1291

6 files changed +170 -33

tests/compile/passes/test_fusion.py

+21 -5

@@ -550,7 +550,9 @@ class TestGatedModel(torch.nn.Module):

550 """Model that uses RMSNormGated + reshape + group FP8 quant + linear.

551

552 Mimics GatedDeltaNetAttention's output projection path where:

553 - - RMSNormGated operates on per-head tensors (N*H, D)

553 + - RMSNormGated operates on per-head tensors, either (N, H, D) as

554 + QwenGatedDeltaNetAttention passes them or flattened to (N*H, D) as

555 + OlmoHybridGatedDeltaNetAttention does

556 - Output is reshaped to (N, H*D) before group quantization + linear

557 """

558

@@ -563,10 +565,12 @@ class TestGatedModel(torch.nn.Module):

565 group_shape: GroupShape,

566 dtype: torch.dtype,

567 use_aiter_quant: bool = True,

568 + flatten_heads: bool = True,

569 ):

570 super().__init__()

571 self.num_heads = num_heads

572 self.head_dim = head_dim

573 + self.flatten_heads = flatten_heads

574 hidden_dim = num_heads * head_dim

575

576 self.norm = RMSNormGated(

@@ -599,8 +603,11 @@ class TestGatedModel(torch.nn.Module):

603 hidden_dim = num_heads * head_dim

604 x = torch.relu(x)

605 z = torch.relu(z)

602 - x_heads = x.reshape(-1, num_heads, head_dim).reshape(-1, head_dim)

603 - z_heads = z.reshape(-1, num_heads, head_dim).reshape(-1, head_dim)

606 + x_heads = x.reshape(-1, num_heads, head_dim)

607 + z_heads = z.reshape(-1, num_heads, head_dim)

608 + if self.flatten_heads:

609 + x_heads = x_heads.reshape(-1, head_dim)

610 + z_heads = z_heads.reshape(-1, head_dim)

611 normed = self.norm(x_heads, z_heads)

612 merged = normed.reshape(-1, hidden_dim)

613 out = self.fp8_linear(merged)

@@ -638,6 +645,8 @@ class _MockGDNLayer:

645 @pytest.mark.parametrize("head_dim", [128])

646 @pytest.mark.parametrize("num_tokens", [8])

647 @pytest.mark.parametrize("eps", [1e-5, 1e-6])

648 +@pytest.mark.parametrize("flatten_heads", [True, False])

649 +@pytest.mark.parametrize("quant_fp8_op", ["-quant_fp8", "+quant_fp8"])

650 @pytest.mark.skipif(

651 (not current_platform.is_rocm() or not IS_AITER_FOUND),

652 reason="Only test on ROCm with aiter package installed",

@@ -648,6 +657,8 @@ def test_aiter_fusion_rmsnorm_gated_quant(

657 head_dim: int,

658 num_tokens: int,

659 eps: float,

660 + flatten_heads: bool,

661 + quant_fp8_op: str,

662 monkeypatch: pytest.MonkeyPatch,

663 ):

664 group_shape = GroupShape(1, 128)

@@ -655,7 +666,7 @@ def test_aiter_fusion_rmsnorm_gated_quant(

666 model_config=ModelConfig(dtype=dtype),

667 compilation_config=CompilationConfig(

668 mode=CompilationMode.VLLM_COMPILE,

658 - custom_ops=["-rms_norm", "-silu_and_mul", "-quant_fp8"],

669 + custom_ops=["-rms_norm", "-silu_and_mul", quant_fp8_op],

670 pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True),

671 ),

672 )

@@ -688,6 +699,7 @@ def test_aiter_fusion_rmsnorm_gated_quant(

699 group_shape=group_shape,

700 dtype=dtype,

701 use_aiter_quant=True,

702 + flatten_heads=flatten_heads,

703 )

704

705 noop_pass = NoOpEliminationPass(vllm_config)

@@ -708,7 +720,11 @@ def test_aiter_fusion_rmsnorm_gated_quant(

720 model_unfused = torch.compile(model, backend=backend2)

721 result_unfused = model_unfused(x, z)

722

711 - torch.testing.assert_close(result_fused, result_unfused, atol=1e-2, rtol=1e-2)

723 + # The opaque AITER quant op preserves a BF16 intermediate in the

724 + # unfused path, while the fused kernel normalizes and quantizes

725 + # directly. Allow for that extra rounding before the FP8 GEMM.

726 + tol = 2e-2 if quant_fp8_op == "+quant_fp8" else 1e-2

727 + torch.testing.assert_close(result_fused, result_unfused, atol=tol, rtol=tol)

728

729 assert fusion_pass.matched_count == 1

730 backend.check_after_ops(model.ops_in_model_after())

tests/kernels/mamba/test_gdn_output_projection.py new

+53

@@ -0,0 +1,53 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Shape contract for ``QwenGatedDeltaNetAttention._output_projection``.

4 +

5 +Every ``forward_*`` already allocates ``core_attn_out`` and ``z`` as

6 +``(N, H, D)``. The helper now norms those tensors as-is and flattens once

7 +for ``out_proj``; it used to be rank-agnostic via ``reshape(z_shape_og)``.

8 +"""

9 +

10 +from __future__ import annotations

11 +

12 +import types

13 +

14 +import pytest

15 +import torch

16 +

17 +from vllm.model_executor.layers.layernorm import RMSNormGated

18 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (

19 + QwenGatedDeltaNetAttention,

20 +)

21 +

22 +

23 +@pytest.mark.parametrize("dtype", [torch.float32])

24 +@torch.inference_mode()

25 +def test_output_projection_norms_per_head_and_flattens(

26 + default_vllm_config,

27 + dtype: torch.dtype,

28 +) -> None:

29 + num_tokens, num_heads, head_dim = 3, 4, 8

30 + hidden = num_heads * head_dim

31 +

32 + layer = types.SimpleNamespace()

33 + layer.norm = RMSNormGated(

34 + head_dim,

35 + eps=1e-5,

36 + group_size=None,

37 + norm_before_gate=True,

38 + device="cpu",

39 + dtype=dtype,

40 + )

41 + layer.out_proj = lambda x: (x, None)

42 + layer._output_projection = types.MethodType(

43 + QwenGatedDeltaNetAttention._output_projection, layer

44 + )

45 +

46 + core_attn_out = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype)

47 + z = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype)

48 + out = layer._output_projection(core_attn_out, z)

49 +

50 + assert core_attn_out.shape == (num_tokens, num_heads, head_dim)

51 + assert z.shape == (num_tokens, num_heads, head_dim)

52 + assert out.shape == (num_tokens, hidden)

53 + assert out.dtype == dtype

tests/kernels/test_fla_layernorm_guard.py

+44

@@ -501,6 +501,50 @@ def test_rmsnorm_gated_forward_native_dtype(

501 torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)

502

503

504 +@pytest.mark.parametrize("dtype", [torch.bfloat16])

505 +@torch.inference_mode()

506 +def test_rmsnorm_gated_3d_matches_flattened_2d(

507 + default_vllm_config,

508 + dtype: torch.dtype,

509 +) -> None:

510 + """Qwen GDN now norms (N, H, D) in place of a flatten to (N*H, D).

511 +

512 + Last-dim RMS + gate must match the flattened layout on both the native

513 + and CUDA/HIP kernels, including a 3D z view that is not contiguous.

514 + """

515 + from vllm.model_executor.layers.layernorm import RMSNormGated

516 +

517 + set_random_seed(42)

518 + device = torch.device(DEVICE)

519 + num_tokens, num_heads, head_dim = 5, 4, 64

520 + layer = RMSNormGated(

521 + head_dim,

522 + eps=1e-5,

523 + group_size=None,

524 + norm_before_gate=True,

525 + device=device,

526 + dtype=dtype,

527 + )

528 +

529 + packed = torch.randn(num_tokens, num_heads, 2, head_dim, dtype=dtype, device=device)

530 + x3 = packed[:, :, 0, :]

531 + z3 = packed[:, :, 1, :]

532 + assert not x3.is_contiguous()

533 + assert not z3.is_contiguous()

534 + x2 = x3.reshape(-1, head_dim).contiguous()

535 + z2 = z3.reshape(-1, head_dim).contiguous()

536 +

537 + native_3d = layer.forward_native(x3, z3)

538 + native_2d = layer.forward_native(x2, z2).reshape(num_tokens, num_heads, head_dim)

539 + torch.testing.assert_close(native_3d, native_2d, atol=0, rtol=0)

540 +

541 + cuda_3d = layer.forward_cuda(x3, z3)

542 + cuda_2d = layer.forward_cuda(x2, z2).reshape(num_tokens, num_heads, head_dim)

543 + torch.testing.assert_close(cuda_3d, cuda_2d, atol=1e-3, rtol=1e-2)

544 + torch.testing.assert_close(cuda_3d, native_3d, atol=1e-2, rtol=1e-2)

545 + assert cuda_3d.shape == (num_tokens, num_heads, head_dim)

546 +

547 +

548 if __name__ == "__main__":

549 # Run a quick smoke test

550 test_layer_norm_fwd_basic(128, 1024, torch.float16, 42, False)

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

+28 -5

@@ -30,6 +30,7 @@ from ..vllm_inductor_pass import (

30 VllmPatternReplacement,

31 _fx_view_to_reshape,

32 fold_consecutive_reshapes,

33 + remove_noop_reshapes,

34 )

35 from .matcher_utils import (

36 MatcherQuantFP8,

@@ -452,9 +453,9 @@ class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):

453 """Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces

454 with rocm_aiter_fused_rms_gated_fp8_group_quant.

455

455 - The norm operates per-head on (N*H, D) tensors. The compiler folds the

456 - reshape chain so after norm the result goes through reshape->merge->quant.

457 - The pattern reshapes from (N*H, D) to (N, H*D) before calling

456 + The norm operates per-head, on either (N*H, D) or (N, H, D). The compiler

457 + folds the reshape chain so after norm the result goes through

458 + reshape->merge->quant. The pattern reshapes to (N, H*D) before calling

459 MatcherQuantFP8 so that _quantize_group_native sees the full hidden dim

460 and computes the correct num_groups.

461 """

@@ -482,6 +483,16 @@ class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):

483 self.head_dim = head_dim

484

485 def register(self, pm_pass: PatternMatcherPass) -> None:

486 + self._register_layout(pm_pass, flatten_heads=True)

487 + # Both layouts reach the norm. With an opaque quant op they trace to

488 + # the same pattern; with a native quant they differ, because its group

489 + # reshape back to (N, H, D) is a no-op the compiler already dropped.

490 + if not self.quant_matcher.enabled:

491 + self._register_layout(pm_pass, flatten_heads=False)

492 +

493 + def _register_layout(

494 + self, pm_pass: PatternMatcherPass, *, flatten_heads: bool

495 + ) -> None:

496 num_heads = self.num_heads

497 head_dim = self.head_dim

498 hidden_dim = num_heads * head_dim

@@ -502,6 +513,11 @@ class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):

513 z: torch.Tensor,

514 weight: torch.Tensor,

515 ) -> tuple[torch.Tensor, torch.Tensor]:

516 + # The pattern matches on op structure rather than rank, so either

517 + # per-head layout can land here. The fused op takes (M, D), and

518 + # both operands are contiguous, so collapsing the heads is a view.

519 + x = x.reshape(-1, head_dim)

520 + z = z.reshape(-1, head_dim)

521 fused = self.FUSED_OP(

522 x=x,

523 weight=weight,

@@ -519,14 +535,21 @@ class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):

535 return fp8_reshaped, scales_reshaped

536

537 n_tokens = 2

522 - x = self.empty(n_tokens * num_heads, head_dim)

523 - z = self.empty(n_tokens * num_heads, head_dim)

538 + shape = (

539 + (n_tokens * num_heads, head_dim)

540 + if flatten_heads

541 + else (n_tokens, num_heads, head_dim)

542 + )

543 + x = self.empty(*shape)

544 + z = self.empty(*shape)

545 w = self.empty(head_dim)

546

547 def trace_fn(*args, **kwargs):

548 gm = pm.fwd_only(*args, **kwargs)

549 _fx_view_to_reshape(gm)

550 fold_consecutive_reshapes(gm)

551 + if not flatten_heads:

552 + remove_noop_reshapes(gm)

553 return gm

554

555 pm.register_replacement(

vllm/compilation/passes/vllm_inductor_pass.py

+21

@@ -277,6 +277,27 @@ def fold_consecutive_reshapes(gm: fx.GraphModule) -> None:

277 gm.graph.erase_node(inp)

278

279

280 +def remove_noop_reshapes(gm: fx.GraphModule) -> None:

281 + """Drop reshape ops whose output shape equals their input shape.

282 +

283 + Companion to :func:`fold_consecutive_reshapes`. ``make_fx`` records a

284 + reshape to the shape the input already has; the compiled graph has already

285 + dropped it, so the pattern only matches once we drop it too.

286 + """

287 + aten_reshape = torch.ops.aten.reshape.default

288 + for node in list(gm.graph.nodes):

289 + if not is_func(node, aten_reshape):

290 + continue

291 + inp = node.args[0]

292 + if not isinstance(inp, fx.Node):

293 + continue

294 + val, inp_val = node.meta.get("val"), inp.meta.get("val")

295 + if val is None or inp_val is None or val.shape != inp_val.shape:

296 + continue

297 + node.replace_all_uses_with(inp)

298 + gm.graph.erase_node(node)

299 +

300 +

301 def _remove_noop_permutes(gm: fx.GraphModule) -> None:

302 for node in gm.graph.nodes:

303 if not is_func(node, torch.ops.aten.permute.default):

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

+3 -23

@@ -852,13 +852,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

852 The RMSNormGated + quant sequence is eligible for fusion

853 by the compilation pass when fuse_norm_quant is enabled.

854 """

855 - z_shape_og = z.shape

856 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])

857 - z = z.reshape(-1, z.shape[-1])

855 core_attn_out = self.norm(core_attn_out, z)

859 - core_attn_out = core_attn_out.reshape(z_shape_og)

860 - core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)

861 - output, _ = self.out_proj(core_attn_out)

856 + output, _ = self.out_proj(core_attn_out.flatten(-2))

857 return output

858

859 def forward_hip(

@@ -1012,15 +1007,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1007 # ============================================================

1008 # Part 3: Output Projection

1009 # ============================================================

1015 - z_shape_og = z.shape

1016 - # Reshape input data into 2D tensor

1017 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])

1018 - z = z.reshape(-1, z.shape[-1])

1019 - core_attn_out = self.norm(core_attn_out, z)

1020 - core_attn_out = core_attn_out.reshape(z_shape_og)

1021 - core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)

1022 - out, _ = self.out_proj(core_attn_out)

1023 - return out

1010 + return self._output_projection(core_attn_out, z)

1011

1012 def forward_cpu(

1013 self,

@@ -1063,14 +1050,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1050 _encode_layer_name(self.prefix),

1051 )

1052

1066 - z_shape_og = z.shape

1067 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])

1068 - z = z.reshape(-1, z.shape[-1])

1069 - core_attn_out = self.norm(core_attn_out, z)

1070 - core_attn_out = core_attn_out.reshape(z_shape_og)

1071 - core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)

1072 - out, _ = self.out_proj(core_attn_out)

1073 - return out

1053 + return self._output_projection(core_attn_out, z)

1054

1055 def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None:

1056 """Warm up GDN prefill kernels during V1 profiling.