feat(mla): add CUDA graph plan update API

csrc/batch_mla_binding.cu

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 #include "batch_mla_config.inc"
+#include "batch_mla_plan_update.cuh"
 #include "tvm/ffi/container/array.h"
 #include "tvm/ffi/container/tuple.h"
 #include "tvm_ffi_utils.h"
@@ -37,3 +38,4 @@ void BatchMLAPagedAttentionRun(TensorView float_workspace_buffer, TensorView int
 
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchMLAPagedAttentionPlan);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, BatchMLAPagedAttentionRun);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(commit_cuda_graph_plan_update, CommitBatchMLACudaGraphPlanUpdate);

csrc/batch_mla_plan_update.cu

@@ -0,0 +1,151 @@
+/*
+ * Copyright (c) 2026 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cuda_runtime.h>
+
+#include <algorithm>
+
+#include "batch_mla_plan_update.cuh"
+
+namespace {
+
+constexpr int kPlanUpdateThreads = 256;
+constexpr int kPlanUpdateBlocks = 128;
+
+void CheckPlanUpdateTensor(TensorView tensor, DLDataType dtype, const char* name) {
+  TVM_FFI_ICHECK_EQ(tensor.device().device_type, kDLCUDA) << name << " must be a CUDA tensor";
+  TVM_FFI_ICHECK(tensor.IsContiguous()) << name << " must be contiguous";
+  TVM_FFI_ICHECK_EQ(tensor.ndim(), 1) << name << " must be a 1D tensor";
+  TVM_FFI_ICHECK_EQ(tensor.dtype(), dtype) << name << " has an invalid dtype";
+}
+
+void CheckSameDevice(TensorView reference, TensorView tensor, const char* name) {
+  TVM_FFI_ICHECK_EQ(tensor.device().device_type, reference.device().device_type)
+      << name << " must be on the live workspace device";
+  TVM_FFI_ICHECK_EQ(tensor.device().device_id, reference.device().device_id)
+      << name << " must be on the live workspace device";
+}
+
+void CheckEqualCapacity(TensorView live, TensorView shadow, const char* name) {
+  TVM_FFI_ICHECK_EQ(shadow.numel(), live.numel())
+      << name << " shadow capacity must match the live tensor";
+}
+
+__global__ void CommitBatchMLACudaGraphPlanUpdateKernel(
+    uint8_t* live_int_workspace, int32_t* live_qo_indptr, int32_t* live_kv_indptr,
+    int32_t* live_kv_indices, int32_t* live_kv_len_arr, const uint8_t* candidate_int_workspace,
+    const int32_t* candidate_qo_indptr, const int32_t* candidate_kv_indptr,
+    const int32_t* source_kv_indices, const int32_t* candidate_kv_len_arr,
+    int64_t staged_int_workspace_bytes, int64_t qo_indptr_elements, int64_t kv_indptr_elements,
+    int64_t kv_len_arr_elements, int64_t live_kv_indices_length, int64_t work_elements) {
+  const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
+  for (int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < work_elements; index += stride) {
+    if (index < staged_int_workspace_bytes) {
+      live_int_workspace[index] = candidate_int_workspace[index];
+    }
+    if (index < qo_indptr_elements) {
+      live_qo_indptr[index] = candidate_qo_indptr[index];
+    }
+    if (index < kv_indptr_elements) {
+      live_kv_indptr[index] = candidate_kv_indptr[index];
+    }
+    if (index < live_kv_indices_length) {
+      live_kv_indices[index] = source_kv_indices[index];
+    }
+    if (index < kv_len_arr_elements) {
+      live_kv_len_arr[index] = candidate_kv_len_arr[index];
+    }
+  }
+}
+
+}  // namespace
+
+void CommitBatchMLACudaGraphPlanUpdate(
+    TensorView live_int_workspace, TensorView live_qo_indptr, TensorView live_kv_indptr,
+    TensorView live_kv_indices, TensorView live_kv_len_arr, TensorView candidate_int_workspace,
+    TensorView candidate_qo_indptr, TensorView candidate_kv_indptr, TensorView source_kv_indices,
+    TensorView candidate_kv_len_arr, int64_t staged_int_workspace_bytes,
+    int64_t live_kv_indices_length) {
+  CheckPlanUpdateTensor(live_int_workspace, dl_uint8, "live_int_workspace");
+  CheckPlanUpdateTensor(live_qo_indptr, dl_int32, "live_qo_indptr");
+  CheckPlanUpdateTensor(live_kv_indptr, dl_int32, "live_kv_indptr");
+  CheckPlanUpdateTensor(live_kv_indices, dl_int32, "live_kv_indices");
+  CheckPlanUpdateTensor(live_kv_len_arr, dl_int32, "live_kv_len_arr");
+  CheckPlanUpdateTensor(candidate_int_workspace, dl_uint8, "candidate_int_workspace");
+  CheckPlanUpdateTensor(candidate_qo_indptr, dl_int32, "candidate_qo_indptr");
+  CheckPlanUpdateTensor(candidate_kv_indptr, dl_int32, "candidate_kv_indptr");
+  CheckPlanUpdateTensor(source_kv_indices, dl_int32, "source_kv_indices");
+  CheckPlanUpdateTensor(candidate_kv_len_arr, dl_int32, "candidate_kv_len_arr");
+
+  CheckSameDevice(live_int_workspace, live_qo_indptr, "live_qo_indptr");
+  CheckSameDevice(live_int_workspace, live_kv_indptr, "live_kv_indptr");
+  CheckSameDevice(live_int_workspace, live_kv_indices, "live_kv_indices");
+  CheckSameDevice(live_int_workspace, live_kv_len_arr, "live_kv_len_arr");
+  CheckSameDevice(live_int_workspace, candidate_int_workspace, "candidate_int_workspace");
+  CheckSameDevice(live_int_workspace, candidate_qo_indptr, "candidate_qo_indptr");
+  CheckSameDevice(live_int_workspace, candidate_kv_indptr, "candidate_kv_indptr");
+  CheckSameDevice(live_int_workspace, source_kv_indices, "source_kv_indices");
+  CheckSameDevice(live_int_workspace, candidate_kv_len_arr, "candidate_kv_len_arr");
+
+  CheckEqualCapacity(live_qo_indptr, candidate_qo_indptr, "qo_indptr");
+  CheckEqualCapacity(live_kv_indptr, candidate_kv_indptr, "kv_indptr");
+  CheckEqualCapacity(live_kv_len_arr, candidate_kv_len_arr, "kv_len_arr");
+  TVM_FFI_ICHECK_GE(staged_int_workspace_bytes, 0)
+      << "staged_int_workspace_bytes must be nonnegative";
+  TVM_FFI_ICHECK_LE(staged_int_workspace_bytes, live_int_workspace.numel())
+      << "staged_int_workspace_bytes exceeds the live workspace capacity";
+  TVM_FFI_ICHECK_LE(staged_int_workspace_bytes, candidate_int_workspace.numel())
+      << "staged_int_workspace_bytes exceeds the candidate workspace capacity";
+  TVM_FFI_ICHECK_GE(live_kv_indices_length, 0) << "live_kv_indices_length must be nonnegative";
+  TVM_FFI_ICHECK_LE(live_kv_indices_length, live_kv_indices.numel())
+      << "live_kv_indices_length exceeds the live index capacity";
+  TVM_FFI_ICHECK_LE(live_kv_indices_length, source_kv_indices.numel())
+      << "live_kv_indices_length exceeds the source index capacity";
+
+  const uintptr_t live_indices_begin = reinterpret_cast<uintptr_t>(live_kv_indices.data_ptr());
+  const uintptr_t live_indices_end = live_indices_begin + live_kv_indices.numel() * sizeof(int32_t);
+  const uintptr_t source_indices_begin = reinterpret_cast<uintptr_t>(source_kv_indices.data_ptr());
+  const uintptr_t source_indices_end =
+      source_indices_begin + source_kv_indices.numel() * sizeof(int32_t);
+  TVM_FFI_ICHECK(source_indices_end <= live_indices_begin ||
+                 live_indices_end <= source_indices_begin)
+      << "source_kv_indices must not overlap live_kv_indices";
+
+  ffi::CUDADeviceGuard device_guard(live_int_workspace.device().device_id);
+  const cudaStream_t stream = get_stream(live_int_workspace.device());
+  const int64_t work_elements =
+      std::max(staged_int_workspace_bytes,
+               std::max(live_qo_indptr.numel(),
+                        std::max(live_kv_indptr.numel(),
+                                 std::max(live_kv_len_arr.numel(), live_kv_indices_length))));
+  CommitBatchMLACudaGraphPlanUpdateKernel<<<kPlanUpdateBlocks, kPlanUpdateThreads, 0, stream>>>(
+      static_cast<uint8_t*>(live_int_workspace.data_ptr()),
+      static_cast<int32_t*>(live_qo_indptr.data_ptr()),
+      static_cast<int32_t*>(live_kv_indptr.data_ptr()),
+      static_cast<int32_t*>(live_kv_indices.data_ptr()),
+      static_cast<int32_t*>(live_kv_len_arr.data_ptr()),
+      static_cast<const uint8_t*>(candidate_int_workspace.data_ptr()),
+      static_cast<const int32_t*>(candidate_qo_indptr.data_ptr()),
+      static_cast<const int32_t*>(candidate_kv_indptr.data_ptr()),
+      static_cast<const int32_t*>(source_kv_indices.data_ptr()),
+      static_cast<const int32_t*>(candidate_kv_len_arr.data_ptr()), staged_int_workspace_bytes,
+      live_qo_indptr.numel(), live_kv_indptr.numel(), live_kv_len_arr.numel(),
+      live_kv_indices_length, work_elements);
+  const cudaError_t error = cudaGetLastError();
+  TVM_FFI_ICHECK_EQ(error, cudaSuccess)
+      << "CommitBatchMLACudaGraphPlanUpdate launch failed: " << cudaGetErrorString(error);
+}

csrc/batch_mla_plan_update.cuh

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2026 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FLASHINFER_BATCH_MLA_PLAN_UPDATE_CUH_
+#define FLASHINFER_BATCH_MLA_PLAN_UPDATE_CUH_
+
+#include <cstdint>
+
+#include "tvm_ffi_utils.h"
+
+void CommitBatchMLACudaGraphPlanUpdate(
+    TensorView live_int_workspace, TensorView live_qo_indptr, TensorView live_kv_indptr,
+    TensorView live_kv_indices, TensorView live_kv_len_arr, TensorView candidate_int_workspace,
+    TensorView candidate_qo_indptr, TensorView candidate_kv_indptr, TensorView source_kv_indices,
+    TensorView candidate_kv_len_arr, int64_t staged_int_workspace_bytes,
+    int64_t live_kv_indices_length);
+
+#endif  // FLASHINFER_BATCH_MLA_PLAN_UPDATE_CUH_

csrc/batch_mla_sm90_binding.cu

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include "batch_mla_plan_update.cuh"
 #include "batch_mla_sm90_config.inc"
 #include "tvm/ffi/container/array.h"
 #include "tvm/ffi/container/tuple.h"
@@ -38,3 +39,4 @@ void BatchMLAPagedAttentionSM90Run(TensorView float_workspace_buffer,
 
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchMLAPagedAttentionSM90Plan);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, BatchMLAPagedAttentionSM90Run);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(commit_cuda_graph_plan_update, CommitBatchMLACudaGraphPlanUpdate);

docs/api/attention.rst

@@ -305,3 +305,62 @@ PageAttention for MLA
     ``DeprecationWarning``.  Explicit ``backend="cutlass"`` callers that omit
     ``plan`` remain supported through a deprecated compatibility adapter when
     ``kv_len`` and ``page_table`` are supplied.
+
+CUDA graph plan updates
+-----------------------
+
+``BatchMLAPagedAttentionWrapper.update_cuda_graph_plan(metadata=...)`` updates
+the dynamic CSR scheduling state of an existing FA2 or FA3 CUDA graph plan
+without changing captured buffer addresses.  Construct the wrapper with
+``use_cuda_graph=True`` and ``enable_cuda_graph_plan_update=True``, provide its
+reserved graph metadata buffers, complete one successful ``plan()``, capture
+``run()``, and call the update outside active CUDA graph capture before replay.
+The update flag is a temporary compatibility opt-in so legacy graph callers do
+not retain update state they never use; it may become the default after the
+legacy private replanning bridge is retired.  The first call that passes the
+wrapper's lifecycle, capability, and capture checks binds the current CUDA
+stream on the wrapper device, even if backend delegation later fails.  Each
+corresponding graph replay must execute on that same stream.  Cross-stream
+replay and concurrent use are unsupported, and the wrapper cannot observe or
+validate the stream used by an external replay.
+
+The update accepts only complete CSR ``MLAPlanMetadata``.  ``qo_indptr``,
+``kv_indptr``, and ``kv_len_arr`` are host control tensors and must be
+contiguous CPU ``torch.int32`` tensors.  ``kv_indices`` must be contiguous
+``torch.int32`` on the wrapper device, must not overlap any capture-reserved
+wrapper buffer, and must remain alive until its queued publication completes.
+Page indices are not read back to the host.  The addressed page-index prefix
+must fit the capture reservation; publication copies only that prefix, so the
+unused reserved tail remains unchanged.
+
+The initial plan freezes the backend and generated module, reserved tensor
+identities and capacities, batch/output shape, layouts, dtypes, scale and LSE
+contracts, launch geometry, ``plan_info``, and staged workspace size.  An
+update that would change any frozen value fails.  FA2 and FA3 retain one device
+candidate schedule, two pinned-host planner/control staging slots, and their
+events during graph planning/warm-up.  When the existing device and pinned
+planner workspaces can hold two staged prefixes, those schedules are disjoint
+views of their unused tails; otherwise FlashInfer allocates only the missing
+region.  The small control tensors remain separate.  No committed schedule or
+page-index image is retained.  An update queries the existing slot events
+without waiting; if both slots are busy, it fails instead of allocating or
+synchronizing.
+CUTLASS, cuTile, and any backend that has not explicitly opted in reject this operation.
+A slot whose event cannot be recorded is poisoned; if both slots are poisoned,
+call ``plan()`` again to create fresh update state.
+
+Validation, planning, staging, and failures before native publication
+submission leave the preceding live plan untouched.  Once publication has been
+submitted, asynchronous CUDA execution or context failures are outside this
+no-sync guarantee by design.
+
+The private ``_cached_module`` and workspace/metadata mirrors used by older
+callers remain behavior-compatible in this release, as do legacy flat/CSR
+``plan()`` forms and the native planner bridge.  This compatibility bridge is
+deprecated in documentation only in this release.  It emits no runtime warning
+because untouched older SGLang accesses ``_cached_module`` and can promote
+``DeprecationWarning`` to an exception.  The public ``plan()`` /
+``update_cuda_graph_plan()`` / ``run()`` lifecycle is the replacement.
+Removing these private compatibility attributes requires a separately
+announced future change and evidence that the applicable support policy no
+longer includes callers that depend on them.