ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update rope_embedding.py

Browse source
This commit is contained in:
Daniel Han 2025-12-12 03:41:09 -08:00
commit 800a6d42e0
1 changed files with 9 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

17
unsloth/kernels/rope_embedding.py
View file

@ -20,13 +20,6 @@ from ..device_type import DEVICE_COUNT
from .utils import calculate_settings, torch_gpu_device, torch_device_stream
@triton.heuristics(
{
"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
"HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
}
)
@triton.jit
def _rope_embedding_QK(
Q,
Q_batch_stride,
@ -104,9 +97,17 @@ def _rope_embedding_QK(
tl.store(k_ptr + half_head_dim + col_offsets, k1 * cos1 + k0 * sin1, mask = mask)
ROPE_GROUP_SIZE: int = 4
_rope_embedding_QK = triton.jit(_rope_embedding_QK)
_rope_embedding_QK = triton.heuristics(
{
"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
"HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
}
)(_rope_embedding_QK)
ROPE_GROUP_SIZE: int = 4
def _rope_embedding(
Q,
Q_row_stride: tl.constexpr,