From 218c5a1f964e0ea79a571450384a26cbaec70042 Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Thu, 3 Apr 2025 10:57:45 +0800
Subject: [PATCH] Update buffer.py

---
 deep_ep/buffer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deep_ep/buffer.py b/deep_ep/buffer.py
index 14948be..831a2e6 100644
--- a/deep_ep/buffer.py
+++ b/deep_ep/buffer.py
@@ -526,7 +526,7 @@ class Buffer:
             out: the in-place output tensor, if set, the kernel will write the result to this tensor and return it directly.
 
         Returns:
-            combined_x: the reduced token tensor, with shape `[num_combined_tokens, num_topk]` and type `torch.bfloat16`.
+            combined_x: the reduced token tensor, with shape `[num_combined_tokens, hidden]` and type `torch.bfloat16`.
             event: the event after executing the kernel (valid only if `async_finish` is set).
             hook: the receiving hook function (valid only if `return_recv_hook` is set).
         """