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 7, 11:30 UTC
[MoE][Perf] Wrap DSV3 QKVAProj GEMM in custom op for torch.compile (#35751)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Mar 2, 2026 at 18:03 UTC
9319044ee9a1afc22327254ab983ead13324942e
1 file changed
+41
-13
vllm/model_executor/models/deepseek_v2.py
+41
-13
@@ -75,6 +75,7 @@ from vllm.model_executor.model_loader.weight_utils import (
75
from vllm.model_executor.models.utils import sequence_parallel_chunk
76
from vllm.platforms import current_platform
77
from vllm.sequence import IntermediateTensors
78
+from vllm.utils.torch_utils import direct_register_custom_op
79
from vllm.v1.attention.backend import AttentionBackend
80
from vllm.v1.attention.backends.mla.indexer import (
81
DeepseekV32IndexerBackend,
@@ -717,6 +718,44 @@ class Indexer(nn.Module):
718
return self.indexer_op(hidden_states, q_fp8, k, weights)
719
720
721
+def _min_latency_fused_qkv_a_proj_impl(
722
+ input_: torch.Tensor,
723
+ weight: torch.Tensor,
724
+) -> torch.Tensor:
725
+ """
726
+ Dynamically run min-latency gemm if num_tokens <= 16.
727
+ This must be wrapped in a custom op because our torch.compile integration
728
+ does not support runtime dispatching on num_tokens.
729
+ """
730
+ num_tokens = input_.shape[0]
731
+ if 0 < num_tokens <= 16:
732
+ output = torch.empty(
733
+ num_tokens,
734
+ weight.shape[0],
735
+ dtype=torch.bfloat16,
736
+ device=input_.device,
737
+ )
738
+ ops.dsv3_fused_a_gemm(output, input_, weight.T)
739
+ return output
740
+ else:
741
+ return torch.nn.functional.linear(input_, weight)
742
+
743
+
744
+def _min_latency_fused_qkv_a_proj_fake(
745
+ input_: torch.Tensor,
746
+ weight: torch.Tensor,
747
+) -> torch.Tensor:
748
+ return input_.new_empty(input_.shape[0], weight.shape[0])
749
+
750
+
751
+direct_register_custom_op(
752
+ op_name="min_latency_fused_qkv_a_proj",
753
+ op_func=_min_latency_fused_qkv_a_proj_impl,
754
+ mutates_args=[],
755
+ fake_impl=_min_latency_fused_qkv_a_proj_fake,
756
+)
757
+
758
+
759
class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
760
def __init__(
761
self,
@@ -752,19 +791,8 @@ class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
791
self,
792
input_,
793
) -> torch.Tensor | tuple[torch.Tensor, torch.nn.Parameter | None]:
755
- num_tokens = input_.shape[0]
756
- if self._use_min_latency_gemm and (0 < num_tokens <= 16):
757
- output = torch.empty(
758
- num_tokens,
759
- 2112,
760
- dtype=torch.bfloat16,
761
- device=input_.device,
762
- )
763
- ops.dsv3_fused_a_gemm(
764
- output,
765
- input_,
766
- self.weight.T,
767
- )
794
+ if self._use_min_latency_gemm:
795
+ output = torch.ops.vllm.min_latency_fused_qkv_a_proj(input_, self.weight)
796
if not self.return_bias:
797
return output
798
output_bias = self.bias if self.skip_bias_add else None