| 23 | |
| 24 | |
| 25 | class IluvatarModelRunner(GPUModelRunner): |
| 26 | def __init__( |
| 27 | self, |
| 28 | fd_config: FDConfig, |
| 29 | device: str, # logic device |
| 30 | device_id: int, # physical device id |
| 31 | rank: int, |
| 32 | local_rank: int, |
| 33 | ): |
| 34 | # Iluvatar does not support cudagraph |
| 35 | fd_config.graph_opt_config.use_cudagraph = False |
| 36 | super(IluvatarModelRunner, self).__init__( |
| 37 | fd_config=fd_config, device=device, device_id=device_id, rank=rank, local_rank=local_rank |
| 38 | ) |
| 39 | assert not self.speculative_decoding, "Iluvatar does not support speculative decoding" |
| 40 | assert self.guided_backend is None, "Iluvatar does not support guided decoding" |
| 41 | assert not self.cache_config.enable_prefix_caching, "Iluvatar does not support prefix caching" |
| 42 | self.mla_cache = envs.FD_ATTENTION_BACKEND == "MLA_ATTN" |
| 43 | assert not self.mla_cache, "Iluvatar does not support MLA" |
| 44 | assert not self.use_cudagraph, "Iluvatar does not support cudagraph" |
| 45 | if self.enable_mm: |
| 46 | assert ( |
| 47 | not self.cache_config.enable_chunked_prefill |
| 48 | ), "Iluvatar does not support chunked prefill for VL model" |
| 49 | # VL neox style = True |
| 50 | emb_shape = self.share_inputs["rope_emb"].shape |
| 51 | if emb_shape[-1] == self.model_config.head_dim // 2: |
| 52 | emb_shape[-1] = self.model_config.head_dim |
| 53 | self.share_inputs["rope_emb"] = paddle.full( |
| 54 | shape=emb_shape, |
| 55 | fill_value=0, |
| 56 | dtype="float32", |
| 57 | ) |
| 58 | |
| 59 | def _initialize_attn_backend(self) -> None: |
| 60 | """ |
| 61 | Initialize attention backends |
| 62 | """ |
| 63 | assert len(self.attn_backends) == 0 |
| 64 | |
| 65 | num_heads = self.model_config.num_attention_heads // self.parallel_config.tensor_parallel_size |
| 66 | self.model_config.kv_num_heads = max( |
| 67 | 1, |
| 68 | int(self.model_config.num_key_value_heads) // self.parallel_config.tensor_parallel_size, |
| 69 | ) |
| 70 | attn_backend = IluvatarAttnBackend( |
| 71 | self.fd_config, |
| 72 | kv_num_heads=self.model_config.kv_num_heads, |
| 73 | num_heads=num_heads, |
| 74 | head_dim=self.model_config.head_dim, |
| 75 | ) |
| 76 | self.attn_backends.append(attn_backend) |