From 67ddac3ea8d5b36978187f92908f555ff71190e4 Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Wed, 23 Sep 2026 23:55:20 -0700 Subject: [PATCH] [CPU][GDN] Support NIXL DS convolution-state layout (#53300) Signed-off-by: Li, Tianmu Co-authored-by: OpenAI Codex Co-authored-by: Li, Jiang --- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 186 ++++++++++++++++++ tests/platforms/test_cpu.py | 75 ++++++- .../layers/mamba/ops/cpu/gdn_attention.py | 121 ++++++++++-- vllm/platforms/cpu.py | 22 ++- 4 files changed, 383 insertions(+), 21 deletions(-) diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index c239f831c6..1ea264239c 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -473,6 +473,20 @@ def _conv_inputs(total_tokens: int): return x, weight, bias +def _conv_states( + num_slots: int, + state_len: int, + dim: int = CONV_DIM, + layout: str = "SD", +) -> torch.Tensor: + if layout == "SD": + return torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16).transpose( + 1, 2 + ) + assert layout == "DS" + return torch.zeros(num_slots, dim, state_len, dtype=torch.bfloat16) + + def _sd_conv_states( num_slots: int, state_len: int, dim: int = CONV_DIM ) -> torch.Tensor: @@ -484,6 +498,39 @@ def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight +@torch.inference_mode() +def test_causal_conv1d_update_cpu_preserves_sd_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """SD state and indices reach the native operator unchanged.""" + num_slots = 2 + state_len = CONV_KERNEL - 1 + conv_state = _sd_conv_states(num_slots, state_len) + state_indices = torch.tensor([1], dtype=torch.int32) + x = torch.zeros(1, CONV_DIM, dtype=torch.bfloat16) + recorded = {} + + def record_update(**kwargs): + recorded.update(kwargs) + return x + + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) + monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", record_update) + out = gdn_attention._causal_conv1d_update_cpu( + x=x, + conv_states=conv_state, + weight=torch.zeros(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16), + bias=None, + silu_activation=False, + conv_state_indices=state_indices, + is_vnni=False, + ) + + assert out is x + assert recorded["conv_states"] is conv_state + assert recorded["conv_state_indices"] is state_indices + + @torch.inference_mode() def test_spec_aware_mixed_routing_preserves_token_order( monkeypatch: pytest.MonkeyPatch, @@ -1123,6 +1170,145 @@ def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False): return out.transpose(0, 1).contiguous(), conv_state +def _python_native_fwd(**kwargs): + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu, + ) + + return causal_conv1d_fn_cpu( + x=kwargs["x"], + weight=kwargs["weight"], + bias=kwargs["bias"], + conv_states=kwargs["conv_states"], + query_start_loc=kwargs["query_start_loc"], + cache_indices=kwargs["cache_indices"], + has_initial_state=kwargs["has_initial_state"], + activation=kwargs["silu_activation"], + ) + + +@torch.inference_mode() +def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + dim = 8 + seq_lens = [2, 3] + state_len = _STATE_LEN + state_indices = torch.tensor([2, 0], dtype=torch.int32) + query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32) + has_initial_state = torch.tensor([False, True]) + x_tokens, weight, bias = _conv_inputs(sum(seq_lens)) + x_tokens = x_tokens[:, :dim] + weight = weight[:dim] + bias = bias[:dim] + x = x_tokens.transpose(0, 1).contiguous() + logical_state = ( + torch.arange(3 * dim * state_len, dtype=torch.float32) + .view(3, dim, state_len) + .to(torch.bfloat16) + ) + state_sd = _conv_states(3, state_len, dim, "SD") + state_ds = _conv_states(3, state_len, dim, "DS") + state_sd.copy_(logical_state) + state_ds.copy_(logical_state) + recorded = [] + + def record_fwd(**kwargs): + recorded.append( + ( + kwargs["conv_states"].shape, + kwargs["conv_states"].stride(), + kwargs["cache_indices"].clone(), + ) + ) + return _python_native_fwd(**kwargs) + + monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_fwd_cpu", record_fwd) + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) + out_sd = gdn_attention._causal_conv1d_fwd_cpu( + x=x, + weight=weight, + bias=bias, + conv_states=state_sd, + query_start_loc=query_start_loc, + cache_indices=state_indices, + has_initial_state=has_initial_state, + silu_activation=True, + is_vnni=False, + ) + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) + out_ds = gdn_attention._causal_conv1d_fwd_cpu( + x=x, + weight=weight, + bias=bias, + conv_states=state_ds, + query_start_loc=query_start_loc, + cache_indices=state_indices, + has_initial_state=has_initial_state, + silu_activation=True, + is_vnni=False, + ) + + torch.testing.assert_close(out_ds, out_sd) + torch.testing.assert_close(state_ds, state_sd) + assert recorded[0][2].equal(state_indices) + assert recorded[1][0] == (2, dim, state_len) + assert recorded[1][1] == (dim * state_len, 1, dim) + assert recorded[1][2].equal(torch.arange(2, dtype=torch.int32)) + + +@torch.inference_mode() +def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + dim = 8 + seq_len = 4 + state_len = seq_len + 5 + state = _conv_states(4, state_len, dim, "DS") + state.copy_( + torch.arange(state.numel(), dtype=torch.float32).view_as(state).to(state.dtype) + ) + original_state = state.clone() + state_indices = torch.tensor([3, 1, 0], dtype=torch.int32) + num_accepted_tokens = torch.tensor([1, 4, 2], dtype=torch.int32) + mutated_scratch = torch.arange( + state_indices.numel() * dim * state_len, dtype=torch.float32 + ).view(state_indices.numel(), dim, state_len) + mutated_scratch = (mutated_scratch + 200).to(state.dtype) + recorded = {} + + def mutate_update(**kwargs): + recorded.update(kwargs) + kwargs["conv_states"].copy_(mutated_scratch) + return torch.zeros_like(kwargs["x"]) + + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) + monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", mutate_update) + gdn_attention._causal_conv1d_update_cpu( + x=torch.zeros(state_indices.numel(), seq_len, dim, dtype=state.dtype), + conv_states=state, + weight=torch.zeros(dim, CONV_KERNEL, dtype=state.dtype), + bias=None, + silu_activation=False, + conv_state_indices=state_indices, + is_vnni=False, + num_accepted_tokens=num_accepted_tokens, + ) + + expected_state = original_state.clone() + expected_state.index_copy_(0, state_indices.to(torch.long), mutated_scratch) + torch.testing.assert_close(state, expected_state) + torch.testing.assert_close( + state.index_select(0, state_indices)[:, :, _STATE_LEN:], + mutated_scratch[:, :, _STATE_LEN:], + ) + torch.testing.assert_close( + recorded["conv_state_indices"], + torch.arange(state_indices.numel(), dtype=torch.int32), + ) + torch.testing.assert_close(recorded["num_accepted_tokens"], num_accepted_tokens) + + @pytest.mark.skipif( not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" ) diff --git a/tests/platforms/test_cpu.py b/tests/platforms/test_cpu.py index fd270315eb..9c7f0b188b 100644 --- a/tests/platforms/test_cpu.py +++ b/tests/platforms/test_cpu.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from types import SimpleNamespace import pytest import torch -from vllm.config import CacheConfig, VllmConfig +from vllm.config import CacheConfig, KVTransferConfig, VllmConfig from vllm.platforms.cpu import CpuPlatform @@ -44,6 +45,7 @@ def _cpu_config( architecture=architecture or model_type, has_inner_state=True, use_mla=False, + hf_config=SimpleNamespace(model_type=model_type), hf_text_config=SimpleNamespace( model_type=model_type, layer_types=layer_types, @@ -150,3 +152,74 @@ def test_cpu_accelerated_gdn_dtype_policy( CpuPlatform.check_and_update_config(config) assert cache_config.mamba_ssm_cache_dtype == expected_dtype + + +@pytest.mark.parametrize( + ( + "connector", + "extra_config", + "explicit_layout", + "avx512_bf16_supported", + "expected_layout", + ), + [ + pytest.param("NixlConnector", None, None, True, "DS", id="nixl"), + pytest.param("NixlPullConnector", None, None, True, "DS", id="nixl-pull"), + pytest.param("NixlPushConnector", None, None, True, "DS", id="nixl-push"), + pytest.param( + "MultiConnector", + { + "connectors": [ + { + "kv_connector": "ExampleConnector", + "kv_connector_extra_config": {}, + }, + { + "kv_connector": "NixlConnector", + "kv_connector_extra_config": {}, + }, + ] + }, + None, + True, + "DS", + id="multi-nixl", + ), + pytest.param("NixlConnector", None, "SD", True, "SD", id="explicit-override"), + pytest.param("OffloadingConnector", None, None, True, "SD", id="non-nixl"), + pytest.param( + "OffloadingConnector", None, None, False, None, id="non-nixl-no-avx" + ), + ], +) +def test_cpu_conv_state_layout_selection( + monkeypatch: pytest.MonkeyPatch, + connector: str, + extra_config: dict | None, + explicit_layout: str | None, + avx512_bf16_supported: bool, + expected_layout: str | None, +) -> None: + layout_env = "VLLM_SSM_CONV_STATE_LAYOUT" + monkeypatch.setattr( + "torch.cpu._is_avx512_bf16_supported", + lambda: avx512_bf16_supported, + ) + + kv_transfer_config = KVTransferConfig( + kv_connector=connector, + kv_connector_extra_config=extra_config or {}, + kv_role="kv_both", + ) + config = _cpu_config( + CacheConfig(mamba_ssm_cache_dtype="float32"), + model_type="qwen3_5", + resolved_dtype="float32", + ) + config.kv_transfer_config = kv_transfer_config + monkeypatch.delenv(layout_env, raising=False) + if explicit_layout is not None: + monkeypatch.setenv(layout_env, explicit_layout) + + CpuPlatform.check_and_update_config(config) + assert os.environ.get(layout_env) == expected_layout diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index b5c1d11a16..f3f3a33ab1 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -27,6 +27,104 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata _CPU_GDN_ATTENTION_OPS_REGISTERED = False +def _stage_ds_conv_states( + conv_states: torch.Tensor, state_indices: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + storage = conv_states.transpose(1, 2).index_select(0, state_indices) + scratch = storage.transpose(1, 2) + identity_indices = torch.arange( + state_indices.numel(), dtype=torch.int32, device=state_indices.device + ) + return scratch, identity_indices + + +def _scatter_ds_conv_states( + conv_states: torch.Tensor, + state_indices: torch.Tensor, + scratch: torch.Tensor, +) -> None: + conv_states.index_copy_(0, state_indices.to(torch.int64), scratch) + + +def _causal_conv1d_fwd_cpu( + *, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + silu_activation: bool, + is_vnni: bool, +) -> torch.Tensor: + if not is_conv_state_dim_first(): + return ops.causal_conv1d_fwd_cpu( + x=x, + weight=weight, + bias=bias, + conv_states=conv_states, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + silu_activation=silu_activation, + is_vnni=is_vnni, + ) + + scratch, identity_indices = _stage_ds_conv_states(conv_states, cache_indices) + out = ops.causal_conv1d_fwd_cpu( + x=x, + weight=weight, + bias=bias, + conv_states=scratch, + query_start_loc=query_start_loc, + cache_indices=identity_indices, + has_initial_state=has_initial_state, + silu_activation=silu_activation, + is_vnni=is_vnni, + ) + _scatter_ds_conv_states(conv_states, cache_indices, scratch) + return out + + +def _causal_conv1d_update_cpu( + *, + x: torch.Tensor, + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + conv_state_indices: torch.Tensor, + is_vnni: bool, + num_accepted_tokens: torch.Tensor | None = None, +) -> torch.Tensor: + if not is_conv_state_dim_first(): + return ops.causal_conv1d_update_cpu( + x=x, + conv_states=conv_states, + weight=weight, + bias=bias, + silu_activation=silu_activation, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + num_accepted_tokens=num_accepted_tokens, + ) + + scratch, identity_indices = _stage_ds_conv_states(conv_states, conv_state_indices) + out = ops.causal_conv1d_update_cpu( + x=x, + conv_states=scratch, + weight=weight, + bias=bias, + silu_activation=silu_activation, + conv_state_indices=identity_indices, + is_vnni=is_vnni, + num_accepted_tokens=num_accepted_tokens, + ) + _scatter_ds_conv_states(conv_states, conv_state_indices, scratch) + return out + + def cpu_gdn_attention_core( mixed_qkv: torch.Tensor, b: torch.Tensor, @@ -109,10 +207,8 @@ def _cpu_gdn_attention_nonspec( conv_state = layer.kv_cache[0] if use_cpp_conv: - # C++ conv requires [num_allocated_slots, kernel - 1, conv_dim] (SD). - if is_conv_state_dim_first(): - raise RuntimeError("C++ CPU GDN attention requires `SD` conv_state layout.") - conv_state = conv_state.transpose(1, 2) + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(1, 2) else: if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) @@ -141,7 +237,7 @@ def _cpu_gdn_attention_nonspec( decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] if use_cpp_conv: - decode_mixed_qkv = ops.causal_conv1d_update_cpu( + decode_mixed_qkv = _causal_conv1d_update_cpu( x=decode_mixed_qkv, conv_states=conv_state, weight=layer.conv1d.weight, @@ -210,7 +306,7 @@ def _cpu_gdn_attention_nonspec( ] if use_cpp_conv: - prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( + prefill_mixed_qkv = _causal_conv1d_fwd_cpu( x=prefill_mixed_qkv.transpose(0, 1), weight=layer.conv1d.weight, bias=layer.conv1d.bias, @@ -416,7 +512,6 @@ def _spec_forward( can_use_native_conv = ( torch.cpu._is_amx_tile_supported() - and not is_conv_state_dim_first() and width == 4 and num_spec_decodes > 0 and bool(torch.all(seq_lens == seq_lens[0]).item()) @@ -424,7 +519,7 @@ def _spec_forward( ) if can_use_native_conv: q_i = int(seq_lens[0].item()) - conv_out = ops.causal_conv1d_update_cpu( + conv_out = _causal_conv1d_update_cpu( x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim), conv_states=conv_buf, weight=layer.conv1d.weight, @@ -509,8 +604,6 @@ def _spec_aware_nonspec( state_indices_tensor = state_indices_tensor.contiguous() is_amx = torch.cpu._is_amx_tile_supported() - if is_amx and is_conv_state_dim_first(): - raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") if not is_amx: conv_weights = _unpacked_conv_weight(layer) @@ -526,7 +619,7 @@ def _spec_aware_nonspec( decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] if is_amx: - decode_mixed_qkv = ops.causal_conv1d_update_cpu( + decode_mixed_qkv = _causal_conv1d_update_cpu( x=decode_mixed_qkv, conv_states=conv_buf, weight=layer.conv1d.weight, @@ -593,7 +686,7 @@ def _spec_aware_nonspec( num_decodes : num_decodes + num_prefills ] if is_amx: - prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( + prefill_mixed_qkv = _causal_conv1d_fwd_cpu( x=prefill_mixed_qkv.transpose(0, 1), weight=layer.conv1d.weight, bias=layer.conv1d.bias, @@ -664,11 +757,9 @@ def _spec_aware_nonspec_subset( prefill_state_indices = prefill_state_indices.contiguous() is_amx = torch.cpu._is_amx_tile_supported() - if is_amx and is_conv_state_dim_first(): - raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") if is_amx: - conv_out = ops.causal_conv1d_fwd_cpu( + conv_out = _causal_conv1d_fwd_cpu( x=mixed_qkv.transpose(0, 1), weight=layer.conv1d.weight, bias=layer.conv1d.bias, diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 359195d4c2..dc7bf0cab5 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -395,11 +395,23 @@ class CpuPlatform(Platform): # Avoid inductor generates num_thread() and breaks the thread binding os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1" - # For efficient conv state memory access. The C++ causal_conv1d - # kernels (VDPBF16PS, no AMX tiles) consume the SD layout on any - # AVX-512BF16 CPU, so apply it beyond AMX (e.g. AMD Zen5/Turin). - if torch.cpu._is_avx512_bf16_supported(): - os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD" + # NIXL's Mamba descriptors require DS conv state storage. Select it + # before cache shapes are created, while preserving an explicit layout. + conv_state_layout_env = "VLLM_SSM_CONV_STATE_LAYOUT" + if conv_state_layout_env not in os.environ: + kv_transfer_config = vllm_config.kv_transfer_config + uses_nixl = kv_transfer_config is not None and any( + kv_transfer_config.has_connector(name) + for name in ( + "NixlConnector", + "NixlPullConnector", + "NixlPushConnector", + ) + ) + if uses_nixl: + os.environ[conv_state_layout_env] = "DS" + elif torch.cpu._is_avx512_bf16_supported(): + os.environ[conv_state_layout_env] = "SD" ld_preload_str = os.getenv("LD_PRELOAD", "") cpu_architecture = Platform.get_cpu_architecture()