Repository: vllm-project/vllm
File: vllm/model_executor/models/deepseek_v2.py (Python)

Code before:
        q = q.view(-1, self.head_dim)
        q_fp8, q_scale = per_token_group_quant_fp8(
            q,
            self.quant_block_size,
            column_major_scales=False,
            use_ue8m0=self.scale_fmt is not None,
        )
        q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
        q_scale = q_scale.view(-1, self.n_head, 1)

        weights, _ = self.weights_proj(hidden_states)
        weights = (
            weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
        )
        weights = weights.squeeze(-1)

        return self.indexer_op(hidden_states, q_fp8, k, weights)


def _min_latency_fused_qkv_a_proj_impl(
    input_: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor:
    """
    Dynamically run min-latency gemm if num_tokens <= 16.
    This must be wrapped in a custom op because torch.compile.
    """
    num_tokens = input_.shape[0]
    if 0 < num_tokens <= 16:
        output = torch.empty(
            num_tokens,
            weight.shape[0],
            dtype=torch.bfloat16,
            device=input_.device,
        )
        ops.dsv3_fused_a_gemm(output, input_, weight.T)
        return output
    else:
        return torch.nn.functional.linear(input_, weight)


def _min_latency_fused_qkv_a_proj_fake(
    input_: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor:
    return input_.new_empty(input_.shape[0], weight.shape[0])


direct_register_custom_op(
    op_name="min_latency_fused_qkv_a_proj",
    op_func=_min_latency_fused_qkv_a_proj_impl,

Review comment (suggestion):
```suggestion
    This must be wrapped in a custom op because our torch.compile integration
    does not support runtime dispatching on num_tokens.
```

Code after:
        q = q.view(-1, self.head_dim)
        q_fp8, q_scale = per_token_group_quant_fp8(
            q,
            self.quant_block_size,
            column_major_scales=False,
            use_ue8m0=self.scale_fmt is not None,
        )
        q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
        q_scale = q_scale.view(-1, self.n_head, 1)

        weights, _ = self.weights_proj(hidden_states)
        weights = (
            weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
        )
        weights = weights.squeeze(-1)

        return self.indexer_op(hidden_states, q_fp8, k, weights)


def _min_latency_fused_qkv_a_proj_impl(
    input_: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor:
    """
    Dynamically run min-latency gemm if num_tokens <= 16.
    This must be wrapped in a custom op because our torch.compile integration
    does not support runtime dispatching on num_tokens.
    """
    num_tokens = input_.shape[0]
    if 0 < num_tokens <= 16:
        output = torch.empty(
            num_tokens,
            weight.shape[0],
            dtype=torch.bfloat16,
            device=input_.device,
        )
        ops.dsv3_fused_a_gemm(output, input_, weight.T)
        return output
    else:
        return torch.nn.functional.linear(input_, weight)


def _min_latency_fused_qkv_a_proj_fake(
    input_: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor:
    return input_.new_empty(input_.shape[0], weight.shape[0])


direct_register_custom_op(
    op_name="min_latency_fused_qkv_a_proj",