Co-authored-by: michaelzhang-ai <michaelzhang-ai@users.noreply.github.com>
This commit is contained in:
@@ -389,6 +389,19 @@ class RotaryEmbedding(MultiPlatformOp):
|
||||
)
|
||||
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:
|
||||
s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
|
||||
s += f", max_position_embeddings={self.max_position_embeddings}"
|
||||
|
||||
Reference in New Issue
Block a user