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:
金黄色葡萄球君君 2026-03-01 15:59:34 +08:00 committed by GitHub
commit 4d3e7d7bcc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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