Add Test

Signed-off-by: Shunkang <[email protected]>

Author: Shunkang

tests/unittest/others/test_kv_cache_transceiver.py

@@ -1,3 +1,5 @@
+import time
+
 import pytest
 import torch
 
@@ -123,3 +125,46 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype,
     assert torch.equal(
         kv_cache_manager_gen.get_buffers(0),
         kv_cache_manager_ctx.get_buffers(0)), "different kv-cache values"
+
+
+@pytest.mark.timeout(120)
+@pytest.mark.parametrize("attention_type",
+                         [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA],
+                         ids=["mha", "mla"])
+def test_cancel_request_in_transmission_ctx(attention_type):
+    # Init kv_cache manager and cache transceiver
+    mapping = Mapping(world_size=1, rank=0)
+    ctx_kv_cache_dtype = DataType.HALF
+    kv_cache_manager_ctx = create_kv_cache_manager(mapping, ctx_kv_cache_dtype)
+
+    cache_transceiver_config = trtllm.CacheTransceiverConfig(
+        backend=trtllm.CacheTransceiverBackendType.DEFAULT,
+        max_tokens_in_buffer=512)
+
+    kv_cache_transceiver_ctx = create_kv_cache_transceiver(
+        mapping, kv_cache_manager_ctx, attention_type, cache_transceiver_config)
+
+    fill_kv_cache_buffer(kv_cache_manager_ctx)
+
+    # init ctx request
+    sampling_params = SamplingParams()
+    ctx_request = LlmRequest(
+        request_id=0,
+        max_new_tokens=1,
+        input_tokens=list(range(256)),
+        sampling_config=tensorrt_llm.bindings.SamplingConfig(
+            sampling_params._get_sampling_config()),
+        is_streaming=False,
+        llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY)
+
+    kv_cache_manager_ctx.impl.add_sequence(ctx_request.py_request_id,
+                                           ctx_request.prompt_len, 1,
+                                           ctx_request)
+    # send ctx request
+    kv_cache_transceiver_ctx.respond_and_send_async(ctx_request)
+
+    time.sleep(10)
+
+    # cancel ctx request
+    is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request)
+    assert is_cancelled