Co-authored-by: michaelzhang-ai <michaelzhang-ai@users.noreply.github.com>
This commit is contained in:
co-authored by
michaelzhang-ai
parent
90a0d66e1e
commit
5e3103a787
@@ -389,6 +389,19 @@ class RotaryEmbedding(MultiPlatformOp):
|
|||||||
)
|
)
|
||||||
return query, key
|
return query, key
|
||||||
|
|
||||||
|
def forward_hip(self, *args, **kwargs):
|
||||||
|
"""HIP/ROCm implementation.
|
||||||
|
|
||||||
|
The JIT kernels (sglang.jit_kernel.pos_enc) used in forward_cuda's
|
||||||
|
fallback path depend on tvm_ffi which invokes nvidia-smi to detect
|
||||||
|
CUDA compute capability. This fails on AMD GPUs, so we use the
|
||||||
|
pure-PyTorch native implementation instead.
|
||||||
|
|
||||||
|
Uses *args/**kwargs because subclasses (MRotaryEmbedding, etc.)
|
||||||
|
have different forward_native() signatures.
|
||||||
|
"""
|
||||||
|
return self.forward_native(*args, **kwargs)
|
||||||
|
|
||||||
def extra_repr(self) -> str:
|
def extra_repr(self) -> str:
|
||||||
s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
|
s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
|
||||||
s += f", max_position_embeddings={self.max_position_embeddings}"
|
s += f", max_position_embeddings={self.max_position_embeddings}"
|
||||||
|
|||||||
Reference in New Issue
Block a user