diff --git a/python/aisimulate/collector/README.md b/python/aisimulate/collector/README.md index 386e1a86e0..8af3d29bb2 100644 --- a/python/aisimulate/collector/README.md +++ b/python/aisimulate/collector/README.md @@ -98,6 +98,14 @@ vLLM 0.31 recurrent qualification uses the same pinned source: Admission permits qualification/collection; publication still requires complete numerical and timing validation of the selected cases. Three-way prefill convolution and speculative verify need their own numerical evidence. + Prefill releases the completed convolution input and its views before + allocating independent core inputs, and releases native output/state/workspace + buffers after every point, including failures. This removes setup overlap; + it does not reduce the [native workspace](https://github.com/vllm-project/FlashKDA/blob/17a037d98da546deb4591e967cf961a43c034d8b/csrc/flash_kda.cpp#L11) + or change callback grids, inputs, dispatch, or timing. Oversized points still + run and retain explicit failure evidence; partial convolution rows do not + make a failed callback complete. CPU lifetime tests cover both prefill paths + and continuation after a core failure. GPU capacity remains a runtime check. - DeepSeek-V4 mHC remains gated on 0.31: its pre rows do not model the fused default serving path, and shared mHC identities omit architecture. A standalone post comparison may use an isolated research artifact, but cannot establish diff --git a/python/aisimulate/collector/vllm/collect_kda.py b/python/aisimulate/collector/vllm/collect_kda.py index 9b54b67621..14b00ef1a0 100644 --- a/python/aisimulate/collector/vllm/collect_kda.py +++ b/python/aisimulate/collector/vllm/collect_kda.py @@ -236,6 +236,7 @@ def run_kda_context_benchmark( for batch_size in batch_size_list: for seq_len in seq_len_list: nt = batch_size * seq_len + flashkda_buffers = {} try: # Resolved at the 0.27.0 era bump: the former FIXME # kernel-limit guard (nt * proj >= 2**31 rejected) is @@ -364,6 +365,8 @@ def run_conv_qkv3(): results["power_stats"], ) + # Conv timing has finished; these views retain an unused 3*nt*proj input. + mixed = q_in = k_in = v_in = None q = torch.randn(1, nt, nh, hd, dtype=dtype, device=device) k = torch.randn(1, nt, nh, hd, dtype=dtype, device=device) v = torch.randn(1, nt, nh, hd, dtype=dtype, device=device) @@ -376,7 +379,6 @@ def run_conv_qkv3(): if prefill_source == "flashkda_fwd": # kda.py:684-691 allocates these before forward; the new # :312-356 API requires them. Do not time allocation. - flashkda_buffers = {} if _check_compat("vllm==0.31.0", vllm_version): workspace_size = torch.ops._flashkda_C.get_workspace_size(nt, nh, batch_size) flashkda_buffers = { @@ -442,6 +444,7 @@ def run_prefill(): cu = idx = has_init = conv_state = q_cs = k_cs = v_cs = None mixed = q_in = k_in = v_in = q = k = v = raw_g = raw_beta = None a_log = dt_bias = init_state = None + flashkda_buffers.clear() _cleanup("context") summary = f"ok={ok} error={err} skip=0" @@ -1429,9 +1432,8 @@ def run_kda_torch( if __name__ == "__main__": import sys - from vllm.version import __version__ as _v - from collector.registry_types import PerfFile + from vllm.version import __version__ as _v print(f"KDA Collector - vLLM {_v}") print(f"SM Version: {get_sm_version()}") diff --git a/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py b/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py index 4f6f053372..1ed78dc938 100644 --- a/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py +++ b/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py @@ -7,6 +7,7 @@ import ast import sys +import weakref from contextlib import contextmanager from pathlib import Path from types import ModuleType, SimpleNamespace @@ -541,6 +542,7 @@ def allocate(*shape, **kwargs): "q": _ContractTensor(1, 128, 12, 128, dtype="bf16"), "init_state": _ContractTensor(2, 12, 128, 128, dtype="fp32"), "flashkda_prefill": lambda *args, **kwargs: calls.append(kwargs), + "flashkda_buffers": {}, } namespace.update(dict.fromkeys(["k", "v", "raw_g", "raw_beta", "a_log", "dt_bias", "cu"], object())) exec(compile(ast.Module(body=branch.body, type_ignores=[]), str(SOURCE_PATH), "exec"), namespace) @@ -556,6 +558,130 @@ def allocate(*shape, **kwargs): assert calls[1]["workspace"] is calls[0]["workspace"] +@pytest.mark.parametrize("backend", ["flashkda_fwd", "chunk_kda_with_fused_gate"]) +@pytest.mark.parametrize("first_core_fails", [False, True]) +def test_prefill_releases_completed_phase_storage_and_failed_point_buffers(monkeypatch, backend, first_core_fails): + from collector.version_resolver import _check_compat + + mixed_refs, buffer_refs, core_calls, conv_calls, timings, published, cleanups = [], [], [], [], [], [], [] + + class Storage: + pass + + class Tensor(_ContractTensor): + def __init__(self, *shape, dtype=None, storage=None, **kwargs): + super().__init__(*shape, dtype=dtype) + self.storage = storage if storage is not None else Storage() + + def transpose(self, first, second): + shape = list(self.shape) + shape[first], shape[second] = shape[second], shape[first] + return Tensor(*shape, dtype=self.dtype, storage=self.storage) + + def split(self, sizes, dim=0): + return [ + Tensor(*self.shape[:dim], size, *self.shape[dim + 1 :], dtype=self.dtype, storage=self.storage) + for size in sizes + ] + + def randn(*shape, **kwargs): + if len(shape) == 2 and shape[1] == 4608: + assert all(ref() is None for ref in buffer_refs), "Previous point retains native buffers" + if len(shape) == 4: + assert all(ref() is None for ref in mixed_refs), "Completed conv storage overlaps core" + tensor = Tensor(*shape, **kwargs) + if len(shape) == 2 and shape[1] == 4608: + mixed_refs.append(weakref.ref(tensor.storage)) + return tensor + + def buffer(*shape, **kwargs): + tensor = Tensor(*shape, **kwargs) + buffer_refs.append(weakref.ref(tensor.storage)) + return tensor + + def native(*args, **kwargs): + core_calls.append(args[0].shape) + if backend == "flashkda_fwd": + assert kwargs["out"].shape == args[0].shape + assert kwargs["final_state"].shape == (1, 12, 128, 128) + assert kwargs["final_state"].dtype == "fp32" + assert kwargs["workspace"].shape == (1234,) + if first_core_fails and len(core_calls) == 1: + raise MemoryError("native core allocation failed") + + def conv(x, weight, bias, **kwargs): + assert x.storage is mixed_refs[-1]() + assert kwargs["metadata"].num_prefills == 1 + conv_calls.append(x.shape) + + def cleanup(tag): + assert all(ref() is None for ref in buffer_refs), "Point cleanup retains native buffers" + cleanups.append(tag) + + for name, attrs in { + "vllm.model_executor.layers.mamba.ops.causal_conv1d": {"causal_conv1d_fn": conv}, + "vllm.models.kimi_k3.nvidia.ops.third_party.kda": {"chunk_kda_with_fused_gate": native}, + "vllm.v1.attention.backends.gdn_attn": {"GDNAttentionMetadata": SimpleNamespace}, + "vllm.v1.attention.backends.utils": {"compute_causal_conv1d_metadata": lambda *a, **kw: ({}, None, None)}, + "vllm.model_executor.layers.mamba.mamba_utils": {"is_conv_state_dim_first": lambda: False}, + }.items(): + module = ModuleType(name) + module.__dict__.update(attrs) + monkeypatch.setitem(sys.modules, name, module) + + @contextmanager + def benchmark(**kwargs): + timings.append((kwargs["num_warmups"], kwargs["num_runs"], kwargs["repeat_n"])) + kwargs["kernel_func"]() + yield {"latency_ms": 1.0, "power_stats": None} + + torch = SimpleNamespace( + device=lambda device: device, + cuda=SimpleNamespace(set_device=lambda device: None), + bfloat16="bf16", + float32="fp32", + int32="int32", + bool="bool", + uint8="uint8", + randn=randn, + zeros=Tensor, + arange=Tensor, + empty=buffer, + empty_like=lambda t: buffer(*t.shape, dtype=t.dtype), + ops=SimpleNamespace(_flashkda_C=SimpleNamespace(get_workspace_size=lambda *args: 1234)), + ) + run = _compile_function( + "run_kda_context_benchmark", + { + "torch": torch, + "_check_compat": _check_compat, + "KDA_LOWER_BOUND": -5.0, + "NUM_WARMUPS": 3, + "NUM_RUNS": 10, + "benchmark_with_power": benchmark, + "_resolve_prefill_kernel": lambda *args: (backend, native), + "_log": lambda common, latency, source, *args: published.append((common.copy(), source)), + "_cleanup": cleanup, + "_format_failures": lambda failures: ";".join(failures), + }, + ) + args = (7168, 4, 12, 128, 12, 128, [1], [2, 4], "moonshotai/Kimi-K3", "unused.txt", "0.31.0") + if first_core_fails: + with pytest.raises(RuntimeError, match="batch_size=1 seq_len=2: MemoryError"): + run(*args) + else: + run(*args) + assert core_calls == [(1, 2, 12, 128), (1, 4, 12, 128)] + assert conv_calls == [(1536, 2)] * 3 + [(1536, 4)] * 3 + assert timings == [(3, 10, 10), (3, 10, 1)] * 2 + assert cleanups == ["context", "context"] + assert [(row["seq_len"], source) for row, source in published] == ( + [(2, "causal_conv1d_fn_qkv3")] + + ([] if first_core_fails else [(2, backend)]) + + [(4, "causal_conv1d_fn_qkv3"), (4, backend)] + ) + + @pytest.mark.parametrize("phase", ["context", "generation"]) def test_kimi031_fixed_profile_gate_admits_only_b200_sd(monkeypatch, phase): namespace = _gate_namespace()