440 lines
15 KiB
Python
440 lines
15 KiB
Python
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_ascend.ops.triton.fla import chunk, chunk_o, chunk_o_update
|
|
from vllm_ascend.utils import enable_custom_op
|
|
|
|
enable_custom_op()
|
|
|
|
|
|
class _FakeKernel:
|
|
def __init__(self):
|
|
self.grid = None
|
|
self.grid_result = None
|
|
self.launch_kwargs: dict[str, object] | None = None
|
|
|
|
def __getitem__(self, grid):
|
|
self.grid = grid
|
|
self.grid_result = grid({"BV": 128})
|
|
|
|
def launch(**kwargs):
|
|
self.launch_kwargs = kwargs
|
|
|
|
return launch
|
|
|
|
|
|
class _DummyTensor:
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
self.shape = (1,)
|
|
self.dtype = torch.float32
|
|
|
|
def unsqueeze(self, dim: int):
|
|
return self
|
|
|
|
def new_empty(self, *shape):
|
|
return _DummyTensor(f"{self.name}.new_empty")
|
|
|
|
def __getitem__(self, item):
|
|
return self
|
|
|
|
def __setitem__(self, item, value):
|
|
return None
|
|
|
|
def __add__(self, other):
|
|
return self
|
|
|
|
def __sub__(self, other):
|
|
return self
|
|
|
|
def transpose(self, dim0, dim1):
|
|
return self
|
|
|
|
def contiguous(self):
|
|
return self
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
|
|
class _GatherResult:
|
|
def __init__(self, items):
|
|
self.items = items
|
|
|
|
def __getitem__(self, item):
|
|
if isinstance(item, tuple):
|
|
item = item[0]
|
|
return self.items[item]
|
|
|
|
|
|
def _patch_missing_cdiv(monkeypatch: pytest.MonkeyPatch, module) -> None:
|
|
if hasattr(module.triton, "cdiv"):
|
|
return
|
|
monkeypatch.setattr(
|
|
module.triton,
|
|
"cdiv",
|
|
lambda x, y: (x + y - 1) // y,
|
|
raising=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("target", ["chunk_o", "chunk_o_update"])
|
|
def test_chunk_leaf_wrappers_use_prebuilt_chunk_offsets(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
target: str,
|
|
):
|
|
fake_kernel = _FakeKernel()
|
|
sentinel = torch.tensor([0, 2, 5], dtype=torch.int32)
|
|
cu_seqlens = torch.tensor([0, 4, 7], dtype=torch.int32)
|
|
|
|
if target == "chunk_o":
|
|
_patch_missing_cdiv(monkeypatch, chunk_o)
|
|
monkeypatch.setattr(chunk_o, "chunk_fwd_kernel_o", fake_kernel)
|
|
monkeypatch.setattr(
|
|
chunk_o,
|
|
"prepare_chunk_offsets",
|
|
lambda *args, **kwargs: pytest.fail("prepare_chunk_offsets should not be called"),
|
|
)
|
|
chunk_o.chunk_fwd_o(
|
|
q=torch.zeros((2, 4, 1, 8), dtype=torch.float32),
|
|
k=torch.zeros((2, 4, 1, 8), dtype=torch.float32),
|
|
v=torch.zeros((2, 4, 1, 16), dtype=torch.float32),
|
|
h=torch.zeros((4, 1, 8, 16), dtype=torch.float32),
|
|
g=torch.zeros((2, 4, 1), dtype=torch.float32),
|
|
cu_seqlens=cu_seqlens,
|
|
chunk_offsets=sentinel,
|
|
)
|
|
else:
|
|
_patch_missing_cdiv(monkeypatch, chunk_o_update)
|
|
monkeypatch.setattr(chunk_o_update, "chunk_fwd_kernel_o_update", fake_kernel)
|
|
monkeypatch.setattr(
|
|
chunk_o_update,
|
|
"prepare_chunk_offsets",
|
|
lambda *args, **kwargs: pytest.fail("prepare_chunk_offsets should not be called"),
|
|
)
|
|
chunk_o_update.chunk_fwd_o_update(
|
|
q=torch.zeros((2, 4, 1, 8), dtype=torch.float32),
|
|
v=torch.zeros((2, 4, 1, 16), dtype=torch.float32),
|
|
h=torch.zeros((4, 1, 8, 16), dtype=torch.float32),
|
|
h_update=torch.zeros((5, 1, 8, 8), dtype=torch.float32),
|
|
updated_h_state=torch.zeros((1, 8, 16), dtype=torch.float32),
|
|
cu_seqlens=cu_seqlens,
|
|
chunk_offsets=sentinel,
|
|
)
|
|
|
|
assert fake_kernel.launch_kwargs is not None
|
|
assert fake_kernel.launch_kwargs["chunk_offsets"] is sentinel
|
|
|
|
|
|
def test_chunk_gated_delta_rule_fwd_threads_prebuilt_chunk_offsets(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
chunk_offsets = torch.tensor([0, 2, 5], dtype=torch.int32)
|
|
update_chunk_offsets = torch.tensor([0, 3, 7], dtype=torch.int32)
|
|
final_chunk_indices = torch.tensor([1, 3], dtype=torch.int32)
|
|
prebuilt_meta = type(
|
|
"PrebuiltMeta",
|
|
(),
|
|
{
|
|
"block_indices_cumsum": None,
|
|
"cu_seqlens_host": (0, 4, 7),
|
|
"chunk_indices_chunk64_host": (0, 0, 1, 0),
|
|
"chunk_indices_chunk64": None,
|
|
"chunk_offsets_chunk64": chunk_offsets,
|
|
"update_chunk_offsets_chunk64": update_chunk_offsets,
|
|
"final_chunk_indices_chunk64": final_chunk_indices,
|
|
"chunk_indices_large_block": None,
|
|
"keep_meta": None,
|
|
"cu_seqlens_kern": None,
|
|
},
|
|
)()
|
|
|
|
q = _DummyTensor("q")
|
|
k = _DummyTensor("k")
|
|
v = _DummyTensor("v")
|
|
g = _DummyTensor("g")
|
|
beta = _DummyTensor("beta")
|
|
initial_state = _DummyTensor("initial_state")
|
|
|
|
non_pcp_calls: list[tuple[str, object]] = []
|
|
pcp_calls: list[tuple[str, object]] = []
|
|
|
|
def run_case(world_size: int, calls: list[tuple[str, object]]):
|
|
group = type(
|
|
"Group",
|
|
(),
|
|
{
|
|
"world_size": world_size,
|
|
"rank_in_group": 0,
|
|
"all_gather": lambda self, value, dim: _GatherResult([_DummyTensor("g0"), _DummyTensor("g1")]),
|
|
},
|
|
)()
|
|
|
|
monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
|
|
monkeypatch.setattr(chunk, "get_pcp_group", lambda: group)
|
|
monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *args, **kwargs: _DummyTensor("g_cumsum"))
|
|
monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *args, **kwargs: _DummyTensor("A"))
|
|
monkeypatch.setattr(chunk, "solve_tril", lambda *args, **kwargs: _DummyTensor("A_solved"))
|
|
monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *args, **kwargs: (_DummyTensor("w"), _DummyTensor("u")))
|
|
monkeypatch.setattr(
|
|
chunk,
|
|
"chunk_gated_delta_rule_fwd_h",
|
|
lambda *args, **kwargs: (_DummyTensor("h"), _DummyTensor("v_new"), _DummyTensor("final_state")),
|
|
)
|
|
monkeypatch.setattr(
|
|
chunk,
|
|
"chunk_gated_delta_rule_fwd_hupdate",
|
|
lambda *args, **kwargs: _DummyTensor("h_update"),
|
|
)
|
|
monkeypatch.setattr(
|
|
chunk.torch,
|
|
"matmul",
|
|
lambda *args, **kwargs: _DummyTensor("matmul"),
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
chunk.torch,
|
|
"zeros_like",
|
|
lambda *args, **kwargs: _DummyTensor("zeros_like"),
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
torch.ops._C_ascend,
|
|
"chunk_gated_delta_rule_fwd_h",
|
|
lambda *args, **kwargs: (_DummyTensor("h"), _DummyTensor("v_new"), _DummyTensor("final_state")),
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
torch.ops._C_ascend,
|
|
"chunk_fwd_o",
|
|
lambda *args, **kwargs: _DummyTensor("o_ascend"),
|
|
raising=False,
|
|
)
|
|
|
|
def fake_chunk_fwd_o(*args, **kwargs):
|
|
calls.append(("o", kwargs["chunk_offsets"]))
|
|
return _DummyTensor("o")
|
|
|
|
def fake_chunk_fwd_o_update(*args, **kwargs):
|
|
calls.append(("o_update", kwargs["chunk_offsets"]))
|
|
return _DummyTensor("h_updated")
|
|
|
|
monkeypatch.setattr(chunk, "chunk_fwd_o", fake_chunk_fwd_o)
|
|
if world_size > 1:
|
|
monkeypatch.setattr(chunk, "chunk_gated_delta_rule_fwd_hupdate", fake_chunk_fwd_o_update)
|
|
|
|
chunk.chunk_gated_delta_rule_fwd(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
g=g,
|
|
beta=beta,
|
|
scale=1.0,
|
|
initial_state=initial_state,
|
|
output_final_state=False,
|
|
cu_seqlens=torch.tensor([0, 4, 7], dtype=torch.int32),
|
|
prebuilt_meta=prebuilt_meta,
|
|
)
|
|
|
|
run_case(1, non_pcp_calls)
|
|
assert non_pcp_calls == []
|
|
|
|
run_case(2, pcp_calls)
|
|
assert pcp_calls == [("o_update", chunk_offsets)]
|
|
|
|
|
|
def test_chunk_gated_delta_rule_fwd_uses_prebuilt_metadata_without_runtime_tolist(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
prebuilt_meta = type(
|
|
"PrebuiltMeta",
|
|
(),
|
|
{
|
|
"block_indices_cumsum": None,
|
|
"cu_seqlens_host": (0, 4, 7),
|
|
"chunk_indices_chunk64_host": (0, 0, 1, 0),
|
|
"chunk_indices_chunk64": torch.tensor([[0, 0], [1, 0]], dtype=torch.int32),
|
|
"chunk_offsets_chunk64": torch.tensor([0, 1, 2], dtype=torch.int32),
|
|
"update_chunk_offsets_chunk64": torch.tensor([0, 2, 4], dtype=torch.int32),
|
|
"final_chunk_indices_chunk64": torch.tensor([1, 3], dtype=torch.int32),
|
|
"chunk_indices_large_block": None,
|
|
"keep_meta": None,
|
|
"cu_seqlens_kern": None,
|
|
},
|
|
)()
|
|
|
|
q = _DummyTensor("q")
|
|
k = _DummyTensor("k")
|
|
v = _DummyTensor("v")
|
|
g = _DummyTensor("g")
|
|
beta = _DummyTensor("beta")
|
|
initial_state = _DummyTensor("initial_state")
|
|
|
|
captured: dict[str, tuple[int, ...] | None] = {}
|
|
|
|
monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
|
|
monkeypatch.setattr(
|
|
chunk,
|
|
"get_pcp_group",
|
|
lambda: type("Group", (), {"world_size": 1, "rank_in_group": 0})(),
|
|
)
|
|
monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *args, **kwargs: _DummyTensor("g_cumsum"))
|
|
monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *args, **kwargs: _DummyTensor("A"))
|
|
monkeypatch.setattr(chunk, "solve_tril", lambda *args, **kwargs: _DummyTensor("A_solved"))
|
|
monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *args, **kwargs: (_DummyTensor("w"), _DummyTensor("u")))
|
|
monkeypatch.setattr(
|
|
torch.ops._C_ascend,
|
|
"chunk_gated_delta_rule_fwd_h",
|
|
lambda *args, **kwargs: (
|
|
captured.update(
|
|
{
|
|
"cu_seqlens": kwargs["cu_seqlens"],
|
|
"chunk_indices": kwargs["chunk_indices"],
|
|
}
|
|
)
|
|
or (_DummyTensor("h"), _DummyTensor("v_new"), _DummyTensor("final_state"))
|
|
),
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
torch.ops._C_ascend,
|
|
"chunk_fwd_o",
|
|
lambda *args, **kwargs: _DummyTensor("o_ascend"),
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
torch.Tensor,
|
|
"tolist",
|
|
lambda self: pytest.fail("runtime should not convert device tensors to host tuples"),
|
|
)
|
|
|
|
chunk.chunk_gated_delta_rule_fwd(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
g=g,
|
|
beta=beta,
|
|
scale=1.0,
|
|
initial_state=initial_state,
|
|
output_final_state=False,
|
|
cu_seqlens=torch.tensor([0, 4, 7], dtype=torch.int32),
|
|
prebuilt_meta=prebuilt_meta,
|
|
)
|
|
|
|
assert captured["cu_seqlens"] == prebuilt_meta.cu_seqlens_host
|
|
assert captured["chunk_indices"] == prebuilt_meta.chunk_indices_chunk64_host
|
|
|
|
|
|
def test_chunk_gated_delta_rule_fwd_pcp_chaining_subtracts_initial_state(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
"""PCP chaining uses (updated_state[i-1] - initial_state), not updated_state[i-1].
|
|
|
|
With s0 != 0 (subsequent prefill chunk), the fix subtracts s0 to avoid
|
|
double-counting Φ_i·s0. Verified by checking the returned final_state
|
|
matches the sequential result Φ_1·(Φ_0·s0+p_0)+p_1.
|
|
"""
|
|
torch.manual_seed(42)
|
|
N, H, K, V = 1, 2, 4, 4
|
|
s0 = torch.randn(N, H, K, V)
|
|
phi_0 = torch.randn(N, H, K, K)
|
|
phi_1 = torch.randn(N, H, K, K)
|
|
p_0 = torch.randn(N, H, K, V)
|
|
p_1 = torch.randn(N, H, K, V)
|
|
|
|
# Each rank computes final_state = Φ_i · s0 + p_i (from shared s0)
|
|
rank0_fs = torch.matmul(phi_0, s0) + p_0
|
|
rank1_fs = torch.matmul(phi_1, s0) + p_1
|
|
# h_update shape [1, N, H, K, K]; after [:, [0], :, :, :] → [1, N, H, K, K]
|
|
h_update_tensor = phi_0.unsqueeze(0)
|
|
|
|
prebuilt_meta = type(
|
|
"PrebuiltMeta",
|
|
(),
|
|
{
|
|
"block_indices_cumsum": None,
|
|
"cu_seqlens_host": (0, N),
|
|
"chunk_indices_chunk64_host": (0, 0),
|
|
"chunk_indices_chunk64": None,
|
|
"chunk_offsets_chunk64": torch.tensor([0, 1], dtype=torch.int32),
|
|
"update_chunk_offsets_chunk64": torch.tensor([0, 2], dtype=torch.int32),
|
|
"final_chunk_indices_chunk64": torch.tensor([0], dtype=torch.int32),
|
|
"chunk_indices_large_block": None,
|
|
"num_decodes": 0,
|
|
"keep_meta": None,
|
|
"cu_seqlens_kern": None,
|
|
},
|
|
)()
|
|
|
|
all_gather_returns = [
|
|
torch.stack([rank0_fs, rank1_fs]), # all_final_state: [2, N, H, K, V]
|
|
torch.stack([phi_0, phi_1]), # all_final_h_update: [2, N, H, K, K]
|
|
]
|
|
|
|
group = type(
|
|
"Group",
|
|
(),
|
|
{
|
|
"world_size": 2,
|
|
"rank_in_group": 0,
|
|
"all_gather": lambda self, value, dim: all_gather_returns.pop(0),
|
|
},
|
|
)()
|
|
|
|
monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
|
|
monkeypatch.setattr(chunk, "get_pcp_group", lambda: group)
|
|
monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *a, **kw: _DummyTensor("g_cumsum"))
|
|
monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *a, **kw: _DummyTensor("A"))
|
|
monkeypatch.setattr(chunk, "solve_tril", lambda *a, **kw: _DummyTensor("A_solved"))
|
|
monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *a, **kw: (_DummyTensor("w"), _DummyTensor("u")))
|
|
monkeypatch.setattr(
|
|
torch.ops._C_ascend,
|
|
"chunk_gated_delta_rule_fwd_h",
|
|
lambda *a, **kw: (_DummyTensor("h"), _DummyTensor("v_new"), rank0_fs),
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
chunk,
|
|
"chunk_gated_delta_rule_fwd_hupdate",
|
|
lambda *a, **kw: h_update_tensor,
|
|
)
|
|
monkeypatch.setattr(
|
|
torch.ops._C_ascend,
|
|
"chunk_fwd_o",
|
|
lambda *a, **kw: _DummyTensor("o_ascendc"),
|
|
raising=False,
|
|
)
|
|
|
|
result = chunk.chunk_gated_delta_rule_fwd(
|
|
q=_DummyTensor("q"),
|
|
k=_DummyTensor("k"),
|
|
v=_DummyTensor("v"),
|
|
g=_DummyTensor("g"),
|
|
beta=_DummyTensor("beta"),
|
|
scale=1.0,
|
|
initial_state=s0,
|
|
output_final_state=False,
|
|
cu_seqlens=torch.tensor([0, N], dtype=torch.int32),
|
|
prebuilt_meta=prebuilt_meta,
|
|
)
|
|
|
|
final_state = result[3]
|
|
# Sequential: Φ_1·(Φ_0·s0 + p_0) + p_1
|
|
expected = torch.matmul(phi_1, torch.matmul(phi_0, s0) + p_0) + p_1
|
|
torch.testing.assert_close(final_state, expected, rtol=1e-4, atol=1e-4)
|