From 4d3e7d7bcc4d511c71f92210fe9612213cd4d043 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=91=E9=BB=84=E8=89=B2=E8=91=A1=E8=90=84=E7=90=83?= =?UTF-8?q?=E5=90=9B=E5=90=9B?= Date: Sun, 1 Mar 2026 15:59:34 +0800 Subject: [PATCH] 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) --- unsloth/kernels/cross_entropy_loss.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index d92229314f..1bb1691776 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -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