| 440 | |
| 441 | |
| 442 | class RowLinear(LinearBase): |
| 443 | |
| 444 | def __init__( |
| 445 | self, |
| 446 | in_features, |
| 447 | out_features, |
| 448 | bias=True, |
| 449 | dtype=None, |
| 450 | tp_group=None, |
| 451 | tp_size=1, |
| 452 | strict_dtype: bool = False, |
| 453 | pad_lda=0, |
| 454 | prefer_managed_weight=True, |
| 455 | is_expert=False, |
| 456 | ): |
| 457 | super().__init__( |
| 458 | local_in_features=in_features // tp_size, |
| 459 | local_out_features=out_features, |
| 460 | bias=bias, |
| 461 | dtype=dtype, |
| 462 | tp_group=tp_group, |
| 463 | tp_size=tp_size, |
| 464 | strict_dtype=strict_dtype, |
| 465 | pad_lda=pad_lda, |
| 466 | prefer_managed_weight=prefer_managed_weight, |
| 467 | ) |
| 468 | |
| 469 | self.tp_dim = 1 |
| 470 | self.tp_size = tp_size |
| 471 | self.is_expert = is_expert |
| 472 | |
| 473 | @classmethod |
| 474 | def tp_split_dim(cls) -> int: |
| 475 | return 1 |
| 476 | |
| 477 | def multiply_collect( |
| 478 | self, |
| 479 | x, |
| 480 | weight, |
| 481 | gemm_plugin: Optional[str] = None, |
| 482 | low_latency_gemm_plugin: Optional[str] = None, |
| 483 | use_fp8: bool = False, |
| 484 | alpha: Optional[np.ndarray] = None, |
| 485 | lora_runtime_params: Optional[LoraRuntimeParams] = None, |
| 486 | lora_hidden_state: Optional[Tensor] = None, |
| 487 | **kwargs): |
| 488 | |
| 489 | gemm_allreduce_plugin = default_net( |
| 490 | ).plugin_config.gemm_allreduce_plugin |
| 491 | if gemm_allreduce_plugin: |
| 492 | if lora_runtime_params != None or lora_hidden_state != None: |
| 493 | raise RuntimeError( |
| 494 | "gemm_allreduce_plugin not supported with lora.") |
| 495 | |
| 496 | output_dtype = self.dtype |
| 497 | if isinstance(output_dtype, str): |
| 498 | output_dtype = str_dtype_to_trt(output_dtype) |
| 499 |
no outgoing calls
no test coverage detected