From 91e7e84ee5a0b482261958edae03c5a752065ee3 Mon Sep 17 00:00:00 2001 From: eeecho <82921722+AliceChenyy@users.noreply.github.com> Date: Tue, 25 Aug 2026 08:42:23 +0800 Subject: [PATCH] [SM120] flash_mla: allocate the page-split buffer outside inference mode (#35116) --- .../kernels/ops/attention/flash_mla_sm120.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index ec43c5508..6f39fa05b 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -400,12 +400,16 @@ def _split_kv_pages_to_64( key = f"flash_mla_sm120_split:{dev}" buf = buffers.get(key) if buf is None or buf.shape[0] < num_dst_pages: - buf = torch.empty( - num_dst_pages, - _BYTES_PER_DST_PAGE_PADDED, - dtype=torch.uint8, - device=dev, - ) + # The first allocation can happen under inference mode (autotune), but + # the buffer is written again during CUDA graph capture outside + # inference mode, where an inference tensor cannot be mutated. + with torch.inference_mode(False): + buf = torch.empty( + num_dst_pages, + _BYTES_PER_DST_PAGE_PADDED, + dtype=torch.uint8, + device=dev, + ) buffers[key] = buf out = buf[:num_dst_pages]