[Bugfix][ROCm] Pre-allocate `wvSplitKrc` static workspaces before KV init (#54247)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Kimi Code <noreply@moonshot.ai>
Nick Hill committed
Aug 28, 2026 at 14:22 UTC
70732942c6e0bcbe83a74cafc9214a0f49fd8929
2 files changed
+37
-6
vllm/model_executor/layers/utils.py
+21
@@ -368,6 +368,27 @@ direct_register_custom_op(
368
)
369
370
371
+@functools.cache
372
+def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None:
373
+ """Eagerly allocate wvSplitKrc's process-lifetime static workspaces.
374
+
375
+ They are otherwise created lazily on the first qualifying GEMM
376
+ (csrc/rocm/skinny_gemms.cu), which can be the first real request — after
377
+ the KV cache backing buffer exists. If one landed in that segment's
378
+ rounding tail, it would pin the entire segment at engine shutdown.
379
+ """
380
+ from vllm.platforms.rocm import on_gfx950
381
+
382
+ if not on_gfx950():
383
+ return
384
+ try:
385
+ x = torch.zeros(16, 1024, dtype=torch.bfloat16, device=device)
386
+ weight = torch.zeros(32, 1024, dtype=torch.bfloat16, device=device)
387
+ ops.wvSplitKrc(x, weight, num_compute_units())
388
+ except Exception:
389
+ logger.debug("wvSplitKrc workspace warmup failed", exc_info=True)
390
+
391
+
392
# Above this weight size, oneDNN's onednn_mm consistently matches or beats
393
# the SGL AMX kernel once M grows past decode-sized batches, and is within
394
# noise of it at decode-sized M -- so larger weights default to oneDNN
vllm/v1/worker/utils.py
+16
-6
@@ -13,6 +13,7 @@ import torch
13
from vllm.config import CacheConfig, VllmConfig
14
from vllm.logger import init_logger
15
from vllm.model_executor.layers.attention import Attention
16
+from vllm.model_executor.layers.utils import warmup_rocm_skinny_gemm_workspaces
17
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
18
from vllm.model_executor.models.utils import extract_layer_index
19
from vllm.platforms import current_platform
@@ -393,12 +394,21 @@ def allocate_kv_cache(
394
sizes = {tensor.size for tensor in kv_cache_config.kv_cache_tensors}
395
assert len(sizes) == 1, "KV cache tensors must share one backing allocation."
396
raw_size = sizes.pop()
396
- page_size = 4096
397
- buf = torch.zeros(
398
- ((raw_size + page_size - 1) // page_size) * page_size,
399
- dtype=torch.int8,
400
- device=device,
401
- )
397
+ # wvSplitKrc's process-lifetime static workspaces (csrc/rocm/skinny_gemms.cu)
398
+ # are created lazily on the first qualifying GEMM. Force that now, before
399
+ # the giant backing allocation below: if one landed in this segment's
400
+ # rounding tail it would pin the whole segment at engine shutdown.
401
+ if current_platform.is_rocm():
402
+ warmup_rocm_skinny_gemm_workspaces(device)
403
+ # Pad to the page granularity MoRIIO needs to register the shared
404
+ # backing as a single RDMA memory region. Other platforms keep the
405
+ # exact-size allocation: NIXL and SimpleCPUOffload rely on
406
+ # storage.nbytes() matching the logical KV size (see #53974).
407
+ page_size = 4096
408
+ buf_size = ((raw_size + page_size - 1) // page_size) * page_size
409
+ else:
410
+ buf_size = raw_size
411
+ buf = torch.zeros(buf_size, dtype=torch.int8, device=device)
412
413
kv_caches: dict[str, torch.Tensor] = {}
414
for tensor in kv_cache_config.kv_cache_tensors: