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