iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CPU][GDN] Support NIXL DS convolution-state layout (#53300)

Browse Source 
Signed-off-by: Li, Tianmu <tianmu.li@intel.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
Tianmu Li
2026-09-24 14:55:20 +08:00
committed by GitHub
co-authored by OpenAI Codex Li, Jiang <jiang1.li@intel.com>
parent 153f2ba1b9
commit 67ddac3ea8
4 changed files with 383 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+186
View File
@@ -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"
)
tests/platforms/test_cpu.py
+74 -1
View File
@@ -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
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+106 -15
View File
@@ -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,
vllm/platforms/cpu.py
+17 -5
View File
@@ -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()