A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 01:30 UTC
Use the same memory for workspace13 and fused_output. (#31531)
Signed-off-by: Andrey Khalyavin <halyavin@yandex-team.ru>
Andrey Khalyavin committed
Jan 18, 2026 at 22:14 UTC
ba29ab441e9bae7d84b9c191ec78256201cda699
1 file changed
+8
-6
vllm/model_executor/layers/fused_moe/modular_kernel.py
+8
-6
@@ -848,15 +848,17 @@ class FusedMoEModularKernel(torch.nn.Module):
848
# We can reuse the memory between cache1 and cache3 because by the
849
# time we need cache3, we're done with cache1.
850
# Construct the entire output that can then be processed in chunks.
851
- # Reuse workspace13 for the output in the non-chunked case as long
852
- # as it is large enough. This will not always be the case for standard
851
+ # Reuse workspace13 for the output in the non-chunked case.
852
+ # This will not always be the case for standard
853
# format experts and with experts that have empty workspaces.
854
- if num_chunks == 1 and prod(workspace13_shape) >= prod(fused_out_shape):
855
- workspace13, workspace2 = current_workspace_manager().get_simultaneous(
856
- (workspace13_shape, workspace_dtype),
854
+ if num_chunks == 1:
855
+ max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
856
+ common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
857
+ ((max_shape_size,), workspace_dtype),
858
(workspace2_shape, workspace_dtype),
859
)
859
- fused_out = _resize_cache(workspace13, fused_out_shape)
860
+ workspace13 = _resize_cache(common_workspace, workspace13_shape)
861
+ fused_out = _resize_cache(common_workspace, fused_out_shape)
862
else:
863
workspace13, workspace2, fused_out = (
864
current_workspace_manager().get_simultaneous(