Helper method to setup GPUModelRunner with different configurations
(self)
| 124 | self.fused_moe = self.setup_fused_moe() |
| 125 | |
| 126 | def setup_model_runner(self): |
| 127 | """Helper method to setup GPUModelRunner with different configurations""" |
| 128 | mock_fd_config = MockFDConfig() |
| 129 | |
| 130 | mock_model_config = MockModelConfig() |
| 131 | mock_cache_config = MockCacheConfig() |
| 132 | |
| 133 | model_runner = GPUModelRunner.__new__(GPUModelRunner) |
| 134 | model_runner.fd_config = mock_fd_config |
| 135 | model_runner.model_config = mock_model_config |
| 136 | model_runner.cache_config = mock_cache_config |
| 137 | model_runner.attn_backends = [MockAttentionBackend()] |
| 138 | model_runner.enable_mm = True |
| 139 | model_runner.cudagraph_only_prefill = False |
| 140 | model_runner.use_cudagraph = False |
| 141 | model_runner.speculative_decoding = False |
| 142 | model_runner.share_inputs = InputBatch(mock_fd_config) |
| 143 | model_runner.share_inputs.init_share_inputs() |
| 144 | model_runner.share_inputs["caches"] = None |
| 145 | model_runner.routing_replay_manager = None |
| 146 | model_runner.exist_prefill_flag = False |
| 147 | |
| 148 | if dist.get_rank() == 0: |
| 149 | model_runner.share_inputs["ids_remove_padding"] = paddle.ones([10]) |
| 150 | else: |
| 151 | model_runner.share_inputs["ids_remove_padding"] = paddle.ones([1]) |
| 152 | |
| 153 | return model_runner |
| 154 | |
| 155 | def setup_fused_moe(self): |
| 156 | mock_fd_config = MockFDConfig() |
no test coverage detected