[II] Reserve Kimi-K3 AttnRes prefill workspace

tests/models/kimi_k3/test_attn_res.py

@@ -194,6 +194,135 @@ def test_attn_res_without_output_norm():
     torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
 
 
+def test_attn_res_reuses_delta_for_output():
+    prefix = torch.randn(17, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    delta = torch.randn_like(prefix)
+    blocks = torch.randn(
+        17, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
+    )
+    norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    qk_weight = (
+        torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
+    )
+    output_norm_weight = torch.ones_like(norm_weight)
+    expected, expected_prefix = _reference(
+        prefix.clone(),
+        delta.clone(),
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        MAX_BLOCKS,
+    )
+    output_pointer = delta.data_ptr()
+
+    actual = attn_res(
+        prefix,
+        delta,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        MAX_BLOCKS,
+        -1,
+        EPS,
+        EPS,
+        output=delta,
+    )
+
+    assert actual.data_ptr() == output_pointer
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+    torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
+
+
+def test_attn_res_uses_an_inactive_workspace_block_for_output():
+    num_tokens = 17
+    prefix = torch.randn(num_tokens, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    block_storage = torch.randn(
+        MAX_BLOCKS,
+        num_tokens,
+        HIDDEN_SIZE,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    blocks = block_storage.permute(1, 0, 2)
+    output = blocks[:, -1]
+    norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    qk_weight = (
+        torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
+    )
+    output_norm_weight = torch.ones_like(norm_weight)
+    expected, expected_prefix = _reference(
+        prefix.clone(),
+        None,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        0,
+    )
+    output_pointer = output.data_ptr()
+
+    actual = attn_res(
+        prefix,
+        None,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        0,
+        0,
+        EPS,
+        EPS,
+        output=output,
+    )
+
+    assert actual.data_ptr() == output_pointer
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+    torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
+    torch.testing.assert_close(blocks[:, 0], expected_prefix, atol=0, rtol=0)
+
+
+def test_attn_res_reused_output_supports_cuda_graph_replay():
+    prefix = torch.randn(17, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    delta = torch.randn_like(prefix)
+    blocks = torch.randn(
+        17, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
+    )
+    norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    qk_weight = (
+        torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
+    )
+    output_norm_weight = torch.ones_like(norm_weight)
+    original_prefix = prefix.clone()
+    original_delta = delta.clone()
+    graph = torch.cuda.CUDAGraph()
+    torch.cuda.synchronize()
+    with torch.cuda.graph(graph):
+        actual = attn_res(
+            prefix,
+            delta,
+            blocks,
+            norm_weight,
+            qk_weight,
+            output_norm_weight,
+            MAX_BLOCKS,
+            -1,
+            EPS,
+            EPS,
+            output=delta,
+        )
+
+    prefix.copy_(original_prefix)
+    delta.copy_(original_delta)
+    graph.replay()
+    first = actual.clone()
+    prefix.copy_(original_prefix)
+    delta.copy_(original_delta)
+    graph.replay()
+    torch.testing.assert_close(actual, first, atol=0, rtol=0)
+
+
 @pytest.mark.parametrize("num_tokens", [0, 1, 17])
 def test_fused_mtp_input(num_tokens: int):
     positions = torch.arange(num_tokens, device="cuda")

tests/models/kimi_k3/test_eagle3.py

@@ -9,15 +9,19 @@
 from vllm.model_executor.models.interfaces import supports_eagle3
 from vllm.models.kimi_k3.nvidia import model as kimi_model
 from vllm.models.kimi_k3.nvidia.model import (
+    KimiDecoderLayer,
     KimiK3ForConditionalGeneration,
     KimiLinearModel,
 )
 
 
 def _make_kimi_linear_model() -> KimiLinearModel:
     model = object.__new__(KimiLinearModel)
+    torch.nn.Module.__init__(model)
+    model.register_buffer("_attn_res_workspace", None, persistent=False)
     object.__setattr__(model, "aux_hidden_state_layers", (2,))
     object.__setattr__(model, "use_sequence_parallel", False)
+    object.__setattr__(model, "reuse_attn_res_output", True)
     return model
 
 
@@ -128,3 +132,176 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch):
     torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
     torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states)
     assert final_attn_res.call_args.args[2] is block_residual
+
+
+def test_kimi_attn_res_workspace_is_reused_and_sliced():
+    model = _make_kimi_linear_model()
+    object.__setattr__(model, "num_attn_res_blocks", 3)
+
+    largest = model._get_attn_res_workspace(torch.empty(8, 4))
+    smaller = model._get_attn_res_workspace(torch.empty(2, 4))
+
+    assert largest.shape == (8, 3, 4)
+    assert smaller.shape == (2, 3, 4)
+    assert smaller.untyped_storage().data_ptr() == largest.untyped_storage().data_ptr()
+
+    grown = model._get_attn_res_workspace(torch.empty(16, 4))
+    assert grown.shape == (16, 3, 4)
+    assert grown.untyped_storage().data_ptr() != largest.untyped_storage().data_ptr()
+
+
+def test_kimi_retained_workspace_releases_cuda_cache(monkeypatch):
+    events = []
+    device = torch.device("cuda", 3)
+    monkeypatch.setattr(
+        torch.accelerator,
+        "synchronize",
+        lambda actual_device: events.append(("synchronize", actual_device)),
+    )
+    monkeypatch.setattr(
+        torch.accelerator,
+        "empty_cache",
+        lambda: events.append(("empty_cache", None)),
+    )
+
+    kimi_model._release_cuda_cache_before_retained_allocation(device)
+
+    assert events == [("synchronize", device), ("empty_cache", None)]
+
+
+def test_kimi_attn_res_workspace_can_be_reserved_before_prefill(monkeypatch):
+    model = _make_kimi_linear_model()
+    object.__setattr__(model, "use_attn_res", True)
+    object.__setattr__(model, "num_attn_res_blocks", 3)
+    object.__setattr__(model, "_max_num_batched_tokens", 16)
+    object.__setattr__(model, "_model_dtype", torch.bfloat16)
+    object.__setattr__(model, "config", SimpleNamespace(hidden_size=4))
+    model.register_parameter(
+        "reference_weight",
+        torch.nn.Parameter(torch.empty(1, dtype=torch.bfloat16)),
+    )
+
+    events = []
+    original_empty = torch.empty
+
+    def record_allocator_release(device):
+        events.append(("release", device.type))
+
+    def record_workspace_allocation(*args, **kwargs):
+        device = kwargs.get("device")
+        if device is not None:
+            events.append(("allocate", device.type))
+        return original_empty(*args, **kwargs)
+
+    monkeypatch.setattr(
+        kimi_model,
+        "_release_cuda_cache_before_retained_allocation",
+        record_allocator_release,
+    )
+    monkeypatch.setattr(torch, "empty", record_workspace_allocation)
+    model.reserve_attn_res_workspace()
+
+    assert events == [("release", "cpu"), ("allocate", "cpu")]
+    workspace = model._attn_res_workspace
+    assert workspace is not None
+    assert workspace.shape == (16, 3, 4)
+    assert workspace.dtype == torch.bfloat16
+    pointer = workspace.untyped_storage().data_ptr()
+    sliced = model._get_attn_res_workspace(torch.empty(8, 4, dtype=torch.bfloat16))
+    assert sliced.untyped_storage().data_ptr() == pointer
+
+
+def _make_attn_res_decoder_layer(*, block_write: bool):
+    layer = object.__new__(KimiDecoderLayer)
+    torch.nn.Module.__init__(layer)
+    object.__setattr__(layer, "use_attn_res", True)
+    object.__setattr__(layer, "reuse_attn_res_output", True)
+    object.__setattr__(layer, "is_block_write_layer", block_write)
+    object.__setattr__(layer, "block_write_idx", 0)
+    object.__setattr__(layer, "prev_valid_blocks", 0)
+    object.__setattr__(
+        layer,
+        "self_attention_res_norm",
+        SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5),
+    )
+    object.__setattr__(
+        layer,
+        "self_attention_res_proj",
+        SimpleNamespace(weight=torch.ones(1, 4)),
+    )
+    object.__setattr__(
+        layer,
+        "input_layernorm",
+        SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5),
+    )
+    object.__setattr__(
+        layer,
+        "mlp_res_norm",
+        SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5),
+    )
+    object.__setattr__(
+        layer,
+        "mlp_res_proj",
+        SimpleNamespace(weight=torch.ones(1, 4)),
+    )
+    object.__setattr__(
+        layer,
+        "post_attention_layernorm",
+        SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5),
+    )
+    return layer
+
+
+def test_kimi_pre_attn_norm_uses_consumed_or_workspace_output(monkeypatch):
+    layer = _make_attn_res_decoder_layer(block_write=True)
+    prefix = torch.randn(2, 4)
+    delta = torch.randn(2, 4)
+    blocks = torch.randn(2, 3, 4)
+    scratch = blocks[:, -1]
+    calls = []
+
+    def record_attn_res(*args, **kwargs):
+        calls.append(kwargs["output"])
+        output = kwargs["output"]
+        return args[0].clone() if output is None else output
+
+    monkeypatch.setattr(kimi_model, "attn_res", record_attn_res)
+
+    first, _, _ = layer._pre_attn_norm(None, blocks, prefix, scratch)
+    later, _, _ = layer._pre_attn_norm(delta, blocks, prefix, scratch)
+
+    assert first is scratch
+    assert later is delta
+    assert calls == [scratch, delta]
+
+    object.__setattr__(layer, "reuse_attn_res_output", False)
+    layer._pre_attn_norm(delta, blocks, prefix, scratch)
+    assert calls[-1] is None
+
+
+def test_kimi_post_attn_norm_reuses_dead_input(monkeypatch):
+    prefix = torch.randn(2, 4)
+    attention_output = torch.randn(2, 4)
+    blocks = torch.randn(2, 3, 4)
+    outputs = []
+
+    def record_attn_res(*args, **kwargs):
+        outputs.append(kwargs["output"])
+        return kwargs["output"]
+
+    monkeypatch.setattr(kimi_model, "attn_res", record_attn_res)
+
+    block_layer = _make_attn_res_decoder_layer(block_write=True)
+    block_hidden, block_prefix, _ = block_layer._post_attn_norm(
+        attention_output, blocks, prefix
+    )
+    regular_layer = _make_attn_res_decoder_layer(block_write=False)
+    regular_hidden, regular_prefix, _ = regular_layer._post_attn_norm(
+        attention_output, blocks, prefix
+    )
+
+    assert block_hidden is prefix
+    assert block_prefix is attention_output
+    assert regular_hidden is attention_output
+    assert regular_prefix is prefix
+    assert outputs == [prefix, attention_output]