| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Model-path tests for fused Qwen3.5 GDN MTP decode.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import types |
| 8 | from typing import cast |
| 9 | from unittest.mock import patch |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | if not (current_platform.is_cuda() and current_platform.has_device_capability(80)): |
| 17 | pytest.skip( |
| 18 | reason="Fused GDN MTP decode requires CUDA compute capability 8.0+.", |
| 19 | allow_module_level=True, |
| 20 | ) |
| 21 | |
| 22 | from tests.v1.attention.utils import ( # noqa: E402 |
| 23 | BatchSpec, |
| 24 | create_common_attn_metadata, |
| 25 | create_vllm_config, |
| 26 | ) |
| 27 | from vllm.config import SpeculativeConfig, set_current_vllm_config # noqa: E402 |
| 28 | from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402 |
| 29 | from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402 |
| 30 | ChunkGatedDeltaRule, |
| 31 | QwenGatedDeltaNetAttention, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402 |
| 34 | MambaStateShapeCalculator, |
| 35 | ) |
| 36 | from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( # noqa: E402 |
| 37 | rmsnorm_fn, |
| 38 | ) |
| 39 | from vllm.utils.torch_utils import _encode_layer_name # noqa: E402 |
| 40 | from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 |
| 41 | GDNAttentionMetadata, |
| 42 | GDNAttentionMetadataBuilder, |
| 43 | ) |
| 44 | from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 |
| 45 | |
| 46 | NUM_SPEC = 3 |
| 47 | SPEC_TOKENS = NUM_SPEC + 1 |
| 48 | H = 1 |
| 49 | HV = 8 |
| 50 | K = 128 |
| 51 | V = 128 |
| 52 | CONV_KERNEL = 4 |
| 53 | CONV_DIM = 2 * H * K + HV * V |
| 54 | BLOCK_SIZE = 16 |
| 55 | PREFIX = "model.layers.0.linear_attn" |
| 56 | EPS = 1e-6 |
| 57 | |
| 58 | |
| 59 | class _TestGatedNorm: |
| 60 | def __init__(self, weight: torch.Tensor, activation: str) -> None: |
| 61 | self.weight = weight |
| 62 | self.bias = None |
| 63 | self.eps = EPS |
| 64 | self.group_size = None |
| 65 | self.norm_before_gate = True |
| 66 | self.activation = activation |
| 67 | |
| 68 | def __call__(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor: |
| 69 | return rmsnorm_fn( |
| 70 | x, |
| 71 | self.weight, |
| 72 | None, |
| 73 | z=z, |
| 74 | eps=self.eps, |
| 75 | norm_before_gate=True, |
| 76 | activation=self.activation, |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | def _make_vllm_config(): |
| 81 | config = create_vllm_config( |
| 82 | model_name="Qwen/Qwen3.5-0.8B", |
| 83 | block_size=BLOCK_SIZE, |
| 84 | hf_config_override={"linear_key_head_dim": K}, |
| 85 | ) |
| 86 | config.additional_config = {"gdn_prefill_backend": "cutedsl"} |
| 87 | config.cache_config.mamba_cache_mode = "none" |
| 88 | config.speculative_config = SpeculativeConfig( |
| 89 | method="ngram", num_speculative_tokens=NUM_SPEC |
| 90 | ) |
| 91 | return config |
| 92 | |
| 93 | |
| 94 | def _build_layer( |
| 95 | vllm_config, |
| 96 | conv_state: torch.Tensor, |
| 97 | ssm_state: torch.Tensor, |
| 98 | a_log: torch.Tensor, |
| 99 | dt_bias: torch.Tensor, |
| 100 | conv_weight: torch.Tensor, |
| 101 | norm_weight: torch.Tensor, |
| 102 | output_gate_activation: str, |
| 103 | ): |
| 104 | layer = types.SimpleNamespace( |
| 105 | prefix=PREFIX, |
| 106 | enable_packed_recurrent_decode=False, |
| 107 | disable_tp_for_ba_proj=False, |
| 108 | tp_size=1, |
| 109 | num_k_heads=H, |
| 110 | num_v_heads=HV, |
| 111 | head_k_dim=K, |
| 112 | head_v_dim=V, |
| 113 | key_dim=K, |
| 114 | value_dim=HV * V, |
| 115 | activation="silu", |
| 116 | A_log=a_log, |
| 117 | dt_bias=dt_bias, |
| 118 | conv1d=types.SimpleNamespace(weight=conv_weight, bias=None), |
| 119 | kv_cache=(conv_state, ssm_state), |
| 120 | norm=_TestGatedNorm(norm_weight, output_gate_activation), |
| 121 | layer_norm_epsilon=EPS, |
| 122 | gdn_decode_kernel="cuda", |
| 123 | ) |
| 124 | with set_current_vllm_config(vllm_config): |
| 125 | layer.chunk_gated_delta_rule = ChunkGatedDeltaRule() |
| 126 | for name in ( |
| 127 | "rearrange_mixed_qkv", |
| 128 | "_forward_core", |
| 129 | "_forward_core_decode_spec_post_conv_fused_norm", |
| 130 | "_forward_core_decode_spec_fused_norm", |
| 131 | "_can_use_fused_gdn_mtp_decode", |
| 132 | "_rms_norm_gated_cuda", |
| 133 | "_forward_core_fused_norm", |
| 134 | "_forward_core_fused_norm_packed", |
| 135 | "split_ba", |
| 136 | ): |
| 137 | setattr( |
| 138 | layer, |
| 139 | name, |
| 140 | types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer), |
| 141 | ) |
| 142 | return layer |
| 143 | |
| 144 | |
| 145 | @pytest.mark.parametrize( |
| 146 | "num_v_heads,expected", |
| 147 | [ |
| 148 | pytest.param(2, True, id="ratio1"), |
| 149 | pytest.param(4, True, id="ratio2"), |
| 150 | pytest.param(6, True, id="ratio3"), |
| 151 | pytest.param(8, True, id="ratio4"), |
| 152 | pytest.param(16, True, id="ratio8"), |
| 153 | pytest.param(5, False, id="non-integral-ratio"), |
| 154 | pytest.param(10, False, id="unsupported-ratio5"), |
| 155 | ], |
| 156 | ) |
| 157 | def test_fused_mtp_head_ratio_guard(num_v_heads: int, expected: bool) -> None: |
| 158 | if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): |
| 159 | pytest.skip("fused GDN decode MTP op is not built") |
| 160 | |
| 161 | layer = types.SimpleNamespace( |
| 162 | num_k_heads=2, |
| 163 | num_v_heads=num_v_heads, |
| 164 | kv_cache=(None, torch.empty(1, dtype=torch.float32, device="cuda")), |
| 165 | gdn_decode_kernel="cuda", |
| 166 | ) |
| 167 | attn_metadata = types.SimpleNamespace( |
| 168 | spec_state_indices_tensor=torch.ones( |
| 169 | 1, SPEC_TOKENS, dtype=torch.int32, device="cuda" |
| 170 | ), |
| 171 | spec_sequence_masks=object(), |
| 172 | num_decodes=0, |
| 173 | num_spec_decodes=1, |
| 174 | ) |
| 175 | |
| 176 | assert ( |
| 177 | QwenGatedDeltaNetAttention._can_use_fused_gdn_mtp_decode( |
| 178 | cast(QwenGatedDeltaNetAttention, layer), |
| 179 | cast(GDNAttentionMetadata, attn_metadata), |
| 180 | ) |
| 181 | is expected |
| 182 | ) |
| 183 | |
| 184 | |
| 185 | @torch.inference_mode() |
| 186 | def test_fused_forward_uses_packed_entrypoint() -> None: |
| 187 | """Fused mode keeps projected QKVZ and BA packed through the model op.""" |
| 188 | device = torch.device("cuda") |
| 189 | num_tokens = 3 |
| 190 | hidden_states = torch.empty(num_tokens, 1, dtype=torch.bfloat16, device=device) |
| 191 | mixed_qkvz = torch.randn( |
| 192 | num_tokens, CONV_DIM + HV * V, dtype=torch.bfloat16, device=device |
| 193 | ) |
| 194 | ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device) |
| 195 | layer = types.SimpleNamespace( |
| 196 | prefix=PREFIX, |
| 197 | enable_fused_gdn_spec_decode=True, |
| 198 | norm=types.SimpleNamespace( |
| 199 | weight=torch.empty(V, dtype=torch.bfloat16, device=device) |
| 200 | ), |
| 201 | num_v_heads=HV, |
| 202 | tp_size=1, |
| 203 | head_v_dim=V, |
| 204 | in_proj_qkvz=lambda _: (mixed_qkvz, None), |
| 205 | in_proj_ba=lambda _: (ba, None), |
| 206 | out_proj=lambda x: (x, None), |
| 207 | ) |
| 208 | layer.forward_cuda = types.MethodType( |
| 209 | QwenGatedDeltaNetAttention.forward_cuda, layer |
| 210 | ) |
| 211 | |
| 212 | def packed_op( |
| 213 | actual_qkvz: torch.Tensor, |
| 214 | actual_ba: torch.Tensor, |
| 215 | output: torch.Tensor, |
| 216 | *, |
| 217 | layer_name: str, |
| 218 | ) -> None: |
| 219 | assert actual_qkvz is mixed_qkvz |
| 220 | assert actual_ba is ba |
| 221 | assert layer_name == _encode_layer_name(PREFIX) |
| 222 | output.fill_(1) |
| 223 | |
| 224 | with ( |
| 225 | patch.object( |
| 226 | torch.ops.vllm, |
| 227 | "qwen_gdn_attention_core_fused_norm_packed", |
| 228 | side_effect=packed_op, |
| 229 | ) as packed_mock, |
| 230 | patch.object(torch.ops.vllm, "qwen_gdn_attention_core") as triton_mock, |
| 231 | ): |
| 232 | output = layer.forward_cuda(hidden_states) |
| 233 | |
| 234 | packed_mock.assert_called_once() |
| 235 | triton_mock.assert_not_called() |
| 236 | torch.testing.assert_close(output, torch.ones_like(output)) |
| 237 | |
| 238 | |
| 239 | @pytest.mark.parametrize( |
| 240 | "seq_lens,query_lens,draft_tokens,expected_fused_calls", |
| 241 | [ |
| 242 | pytest.param([128], [SPEC_TOKENS], [NUM_SPEC], 1, id="pure-mtp"), |
| 243 | pytest.param( |
| 244 | [128, 96], |
| 245 | [SPEC_TOKENS, 64], |
| 246 | [NUM_SPEC, -1], |
| 247 | 0, |
| 248 | id="mixed-mtp-falls-back", |
| 249 | ), |
| 250 | pytest.param([96], [64], [-1], 0, id="pure-prefill"), |
| 251 | pytest.param([128], [1], [-1], 0, id="pure-decode"), |
| 252 | ], |
| 253 | ) |
| 254 | @pytest.mark.parametrize("output_gate_activation", ["silu", "sigmoid"]) |
| 255 | @torch.inference_mode() |
| 256 | def test_fused_model_path_matches_reference( |
| 257 | seq_lens: list[int], |
| 258 | query_lens: list[int], |
| 259 | draft_tokens: list[int], |
| 260 | expected_fused_calls: int, |
| 261 | output_gate_activation: str, |
| 262 | ) -> None: |
| 263 | """Fused MTP and its mixed/prefill/decode fallbacks match the reference.""" |
| 264 | torch.manual_seed(1) |
| 265 | device = torch.device("cuda") |
| 266 | vllm_config = _make_vllm_config() |
| 267 | builder = GDNAttentionMetadataBuilder( |
| 268 | kv_cache_spec=MambaSpec( |
| 269 | block_size=BLOCK_SIZE, |
| 270 | shapes=((16, 64),), |
| 271 | dtypes=(torch.float16,), |
| 272 | num_speculative_blocks=NUM_SPEC, |
| 273 | ), |
| 274 | layer_names=[PREFIX], |
| 275 | vllm_config=vllm_config, |
| 276 | device=device, |
| 277 | ) |
| 278 | batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) |
| 279 | common = create_common_attn_metadata( |
| 280 | batch, BLOCK_SIZE, device, arange_block_indices=True |
| 281 | ) |
| 282 | common.block_table_tensor.add_(1) |
| 283 | with set_current_vllm_config(vllm_config): |
| 284 | metadata = builder.build( |
| 285 | common_prefix_len=0, |
| 286 | common_attn_metadata=common, |
| 287 | num_accepted_tokens=torch.ones( |
| 288 | batch.batch_size, dtype=torch.int32, device=device |
| 289 | ), |
| 290 | num_decode_draft_tokens_cpu=torch.tensor(draft_tokens, dtype=torch.int32), |
| 291 | ) |
| 292 | |
| 293 | state_indices = [ |
| 294 | indices |
| 295 | for indices in ( |
| 296 | metadata.spec_state_indices_tensor, |
| 297 | metadata.non_spec_state_indices_tensor, |
| 298 | ) |
| 299 | if indices is not None and indices.numel() > 0 |
| 300 | ] |
| 301 | pool_size = max(int(indices.max().item()) for indices in state_indices) + 1 |
| 302 | conv_state_shape, temporal_state_shape = ( |
| 303 | MambaStateShapeCalculator.gated_delta_net_state_shape( |
| 304 | 1, H, HV, K, V, CONV_KERNEL, NUM_SPEC |
| 305 | ) |
| 306 | ) |
| 307 | conv_state_seed = 0.05 * torch.randn( |
| 308 | pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device |
| 309 | ) |
| 310 | ssm_state_seed = 0.01 * torch.randn( |
| 311 | pool_size, *temporal_state_shape, dtype=torch.float32, device=device |
| 312 | ) |
| 313 | a_log = 0.1 * torch.randn(HV, dtype=torch.float32, device=device) |
| 314 | dt_bias = 0.1 * torch.randn(HV, dtype=torch.float32, device=device) |
| 315 | conv_weight = 0.1 * torch.randn( |
| 316 | CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device |
| 317 | ) |
| 318 | norm_weight = torch.randn(V, dtype=torch.float32, device=device) |
| 319 | num_tokens = batch.compute_num_tokens() |
| 320 | mixed_qkv = 0.1 * torch.randn( |
| 321 | num_tokens, CONV_DIM, dtype=torch.bfloat16, device=device |
| 322 | ) |
| 323 | b = 0.1 * torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) |
| 324 | a = 0.1 * torch.randn_like(b) |
| 325 | output_gate = 0.1 * torch.randn( |
| 326 | num_tokens, HV, V, dtype=torch.bfloat16, device=device |
| 327 | ) |
| 328 | mixed_qkvz = torch.cat((mixed_qkv, output_gate.flatten(1)), dim=-1) |
| 329 | ba = torch.cat((b, a), dim=-1) |
| 330 | context = types.SimpleNamespace(attn_metadata={PREFIX: metadata}) |
| 331 | |
| 332 | reference_layer = _build_layer( |
| 333 | vllm_config, |
| 334 | conv_state_seed.clone(), |
| 335 | ssm_state_seed.clone(), |
| 336 | a_log, |
| 337 | dt_bias, |
| 338 | conv_weight, |
| 339 | norm_weight, |
| 340 | output_gate_activation, |
| 341 | ) |
| 342 | reference_out = torch.zeros_like(output_gate) |
| 343 | with patch.object( |
| 344 | qwen_gdn_linear_attn, "get_forward_context", return_value=context |
| 345 | ): |
| 346 | reference_layer._forward_core( |
| 347 | mixed_qkv=mixed_qkv.clone(), |
| 348 | b=b, |
| 349 | a=a, |
| 350 | core_attn_out=reference_out, |
| 351 | ) |
| 352 | reference_out = reference_layer.norm(reference_out, output_gate) |
| 353 | |
| 354 | fused_layer = _build_layer( |
| 355 | vllm_config, |
| 356 | conv_state_seed.clone(), |
| 357 | ssm_state_seed.clone(), |
| 358 | a_log, |
| 359 | dt_bias, |
| 360 | conv_weight, |
| 361 | norm_weight, |
| 362 | output_gate_activation, |
| 363 | ) |
| 364 | context.no_compile_layers = {PREFIX: fused_layer} |
| 365 | fused_out = torch.zeros_like(output_gate) |
| 366 | fused_op = qwen_gdn_linear_attn.ops.fused_gdn_decode_post_conv_mtp |
| 367 | with ( |
| 368 | patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=context), |
| 369 | patch.object( |
| 370 | qwen_gdn_linear_attn.ops, |
| 371 | "fused_gdn_decode_post_conv_mtp", |
| 372 | wraps=fused_op, |
| 373 | ) as fused_mock, |
| 374 | ): |
| 375 | torch.ops.vllm.qwen_gdn_attention_core_fused_norm_packed( |
| 376 | mixed_qkvz.clone(), |
| 377 | ba, |
| 378 | fused_out, |
| 379 | layer_name=_encode_layer_name(PREFIX), |
| 380 | ) |
| 381 | |
| 382 | assert fused_mock.call_count == expected_fused_calls |
| 383 | torch.testing.assert_close( |
| 384 | fused_layer.kv_cache[0], reference_layer.kv_cache[0], atol=0, rtol=0 |
| 385 | ) |
| 386 | torch.testing.assert_close(fused_out, reference_out, atol=3e-2, rtol=3e-2) |
| 387 | torch.testing.assert_close( |
| 388 | fused_layer.kv_cache[1], |
| 389 | reference_layer.kv_cache[1], |
| 390 | atol=3e-2, |
| 391 | rtol=3e-2, |
| 392 | ) |