perf(ROCm): optimize chunked CE loss num_warps for RDNA GPUs (#4123)
Use 16 warps for RDNA in the chunked cross-entropy forward kernel
(large vocab > 65536), matching the existing CDNA optimization.
Benchmarked on W7900 (gfx1100) with actual unsloth kernels (5 trials, median):
- Chunked CE forward (BS=65536): 16 warps = 2.4-2.6x faster than 32
- All other kernels (LayerNorm, RoPE, SwiGLU): default heuristic is
already optimal for RDNA; no modification needed.
Depends on: #4109 (provides is_rdna() detection)
This commit is contained in:
parent
0cc694158f
commit
4d3e7d7bcc
1 changed files with 2 additions and 1 deletions
|
|
@ -22,6 +22,7 @@ from .utils import (
|
|||
triton_cast,
|
||||
torch_gpu_device,
|
||||
is_cdna,
|
||||
is_rdna,
|
||||
)
|
||||
from transformers.models.llama.modeling_llama import logger
|
||||
from unsloth_zoo.utils import Version
|
||||
|
|
@ -364,7 +365,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
SOFTCAP = logit_softcapping,
|
||||
DO_LOGIT_SCALING = DO_LOGIT_SCALING,
|
||||
LOGIT_SCALE = logit_scaling,
|
||||
num_warps = 32 if not is_cdna() else 16,
|
||||
num_warps = 16 if is_cdna() or is_rdna() else 32,
|
||||
)
|
||||
# logsumexp(chunked_logsumexp) - x
|
||||
# Do the -x separately
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue