Update buffer.py

This commit is contained in:
fzyzcjy 2025-04-03 10:57:45 +08:00 committed by GitHub
parent 26fa72d80f
commit 218c5a1f96
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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).
"""