gemini-code-assist[bot] commented on code in PR #19904:
URL: https://github.com/apache/tvm/pull/19904#discussion_r3489147698
##########
src/runtime/vm/paged_kv_cache.cc:
##########
@@ -2289,7 +2289,19 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
}
// - Sync Tensors to GPU.
SyncAuxArrayToDevice();
+ // FlashInfer's plan kernels no longer take an explicit stream argument;
they
+ // run on the device's *current* stream. Make the copy stream current
around
+ // the plan so its workspace writes happen on the copy stream -- matching
the
+ // aux-array copies above and the copy->compute synchronization below (this
+ // preserves the prior behavior where copy_stream_ was passed explicitly).
+ bool plan_on_copy_stream = copy_stream_ != nullptr && copy_stream_ !=
compute_stream_;
+ if (plan_on_copy_stream) {
+ DeviceAPI::Get(device_)->SetStream(device_, copy_stream_);
+ }
KernelBeginForward();
+ if (plan_on_copy_stream) {
+ DeviceAPI::Get(device_)->SetStream(device_, compute_stream_);
+ }
Review Comment:
If `KernelBeginForward()` throws an exception, the stream will not be
restored to `compute_stream_`, leaving the device's current stream corrupted as
`copy_stream_`. This can cause subsequent operations to run on the wrong
stream, leading to silent data corruption or race conditions. Using a local
RAII guard ensures that the stream is always restored correctly even in the
presence of exceptions.
```suggestion
bool plan_on_copy_stream = copy_stream_ != nullptr && copy_stream_ !=
compute_stream_;
struct StreamGuard {
Device device;
TVMStreamHandle restore_stream;
bool active;
~StreamGuard() {
if (active) {
DeviceAPI::Get(device)->SetStream(device, restore_stream);
}
}
} guard{device_, compute_stream_, plan_on_copy_stream};
if (plan_on_copy_stream) {
DeviceAPI::Get(device_)->SetStream(device_, copy_stream_);
}
KernelBeginForward();
```
##########
src/runtime/vm/attn_backend.h:
##########
@@ -49,6 +49,134 @@ enum class AttnBackendKind : int {
kFlashInfer = 1,
};
+/*!
+ * \brief Return a zero-copy alias of \p t whose `byte_offset` is folded into
the
+ * data pointer, so the resulting tensor has `byte_offset == 0`.
+ *
+ * FlashInfer 0.6.3 kernels read tensors from `data` directly and do NOT honor
+ * the DLPack `byte_offset` field. mlc's auxiliary index tensors (qo_indptr,
+ * kv_indptr, page_indptr, page_indices, length_info, ...) are views packed
into
+ * a shared workspace and therefore carry a non-zero `byte_offset`. Passing
them
+ * as-is makes FlashInfer read the wrong addresses; this helper rebases them.
+ */
+inline ffi::Tensor ZeroByteOffsetView(const Tensor& t) {
+ if (t->byte_offset == 0) return t;
+ auto* holder = new Tensor(t); // keep the underlying storage alive
+ auto* managed = new DLManagedTensor();
+ managed->manager_ctx = holder;
+ managed->deleter = [](DLManagedTensor* self) {
+ delete[] self->dl_tensor.shape;
+ delete[] self->dl_tensor.strides;
+ delete static_cast<Tensor*>(self->manager_ctx);
+ delete self;
+ };
+ DLTensor& dl = managed->dl_tensor;
+ dl.data = static_cast<void*>(static_cast<char*>(t->data) + t->byte_offset);
+ dl.device = t->device;
+ dl.ndim = t->ndim;
+ dl.dtype = t->dtype;
+ dl.shape = new int64_t[t->ndim];
+ dl.strides = nullptr;
+ for (int i = 0; i < t->ndim; ++i) dl.shape[i] = t->shape[i];
+ if (t->strides != nullptr) {
+ dl.strides = new int64_t[t->ndim];
+ for (int i = 0; i < t->ndim; ++i) dl.strides[i] = t->strides[i];
+ }
+ dl.byte_offset = 0;
+ return tvm::ffi::Tensor::FromDLPack(managed, /*require_alignment=*/0,
+ /*require_contiguous=*/false);
+}
Review Comment:
The current implementation of `ZeroByteOffsetView` is not exception-safe. If
any allocation (such as `new DLManagedTensor()` or `new int64_t[]`) throws
`std::bad_alloc`, or if `tvm::ffi::Tensor::FromDLPack` throws an exception,
previously allocated raw pointers (like `holder` or `managed`) will leak. Using
`std::unique_ptr` and a `try-catch` block ensures that all allocated resources
are safely cleaned up if an exception occurs.
```c
inline ffi::Tensor ZeroByteOffsetView(const Tensor& t) {
if (t->byte_offset == 0) return t;
auto holder = std::make_unique<Tensor>(t);
auto managed = std::make_unique<DLManagedTensor>();
DLTensor& dl = managed->dl_tensor;
dl.data = static_cast<void*>(static_cast<char*>(t->data) + t->byte_offset);
dl.device = t->device;
dl.ndim = t->ndim;
dl.dtype = t->dtype;
auto shape = std::make_unique<int64_t[]>(t->ndim);
for (int i = 0; i < t->ndim; ++i) shape[i] = t->shape[i];
std::unique_ptr<int64_t[]> strides = nullptr;
if (t->strides != nullptr) {
strides = std::make_unique<int64_t[]>(t->ndim);
for (int i = 0; i < t->ndim; ++i) strides[i] = t->strides[i];
}
dl.shape = shape.release();
dl.strides = strides.release();
dl.byte_offset = 0;
managed->manager_ctx = holder.release();
managed->deleter = [](DLManagedTensor* self) {
delete[] self->dl_tensor.shape;
delete[] self->dl_tensor.strides;
delete static_cast<Tensor*>(self->manager_ctx);
delete self;
};
DLManagedTensor* raw_managed = managed.release();
try {
return tvm::ffi::Tensor::FromDLPack(raw_managed, /*require_alignment=*/0,
/*require_contiguous=*/false);
} catch (...) {
raw_managed->deleter(raw_managed);
throw;
}
}
```
##########
src/runtime/vm/attn_backend.h:
##########
@@ -49,6 +49,134 @@ enum class AttnBackendKind : int {
kFlashInfer = 1,
};
+/*!
+ * \brief Return a zero-copy alias of \p t whose `byte_offset` is folded into
the
+ * data pointer, so the resulting tensor has `byte_offset == 0`.
+ *
+ * FlashInfer 0.6.3 kernels read tensors from `data` directly and do NOT honor
+ * the DLPack `byte_offset` field. mlc's auxiliary index tensors (qo_indptr,
+ * kv_indptr, page_indptr, page_indices, length_info, ...) are views packed
into
+ * a shared workspace and therefore carry a non-zero `byte_offset`. Passing
them
+ * as-is makes FlashInfer read the wrong addresses; this helper rebases them.
+ */
+inline ffi::Tensor ZeroByteOffsetView(const Tensor& t) {
+ if (t->byte_offset == 0) return t;
+ auto* holder = new Tensor(t); // keep the underlying storage alive
+ auto* managed = new DLManagedTensor();
+ managed->manager_ctx = holder;
+ managed->deleter = [](DLManagedTensor* self) {
+ delete[] self->dl_tensor.shape;
+ delete[] self->dl_tensor.strides;
+ delete static_cast<Tensor*>(self->manager_ctx);
+ delete self;
+ };
+ DLTensor& dl = managed->dl_tensor;
+ dl.data = static_cast<void*>(static_cast<char*>(t->data) + t->byte_offset);
+ dl.device = t->device;
+ dl.ndim = t->ndim;
+ dl.dtype = t->dtype;
+ dl.shape = new int64_t[t->ndim];
+ dl.strides = nullptr;
+ for (int i = 0; i < t->ndim; ++i) dl.shape[i] = t->shape[i];
+ if (t->strides != nullptr) {
+ dl.strides = new int64_t[t->ndim];
+ for (int i = 0; i < t->ndim; ++i) dl.strides[i] = t->strides[i];
+ }
+ dl.byte_offset = 0;
+ return tvm::ffi::Tensor::FromDLPack(managed, /*require_alignment=*/0,
+ /*require_contiguous=*/false);
+}
+
+/*!
+ * \brief Build a strided, zero-copy view selecting the key (which=0) or value
+ * (which=1) sub-tensor from a combined paged KV tensor of shape
+ * (num_pages, 2, num_heads, page_size, head_dim), yielding a
+ * (num_pages, num_heads, page_size, head_dim) tensor that shares storage with
+ * `pages`. FlashInfer 0.6.3 takes separate key/value paged caches and reads
the
+ * tensor strides, so a strided view avoids an explicit split/copy.
+ */
+inline ffi::Tensor PagedKVCacheView(const Tensor& pages, int64_t which) {
+ TVM_FFI_ICHECK_EQ(pages->ndim, 5);
+ TVM_FFI_ICHECK_EQ(pages->shape[1], 2);
+ int64_t num_pages = pages->shape[0];
+ int64_t num_heads = pages->shape[2];
+ int64_t page_size = pages->shape[3];
+ int64_t head_dim = pages->shape[4];
+ int64_t inner = num_heads * page_size * head_dim;
+ int64_t elem_bytes = (pages->dtype.bits * pages->dtype.lanes + 7) / 8;
+
+ auto* holder = new Tensor(pages); // keep the underlying storage alive
+ auto* managed = new DLManagedTensor();
+ managed->manager_ctx = holder;
+ managed->deleter = [](DLManagedTensor* self) {
+ delete[] self->dl_tensor.shape;
+ delete[] self->dl_tensor.strides;
+ delete static_cast<Tensor*>(self->manager_ctx);
+ delete self;
+ };
+ DLTensor& dl = managed->dl_tensor;
+ dl.data = static_cast<void*>(static_cast<char*>(pages->data) +
pages->byte_offset +
+ which * inner * elem_bytes);
+ dl.device = pages->device;
+ dl.ndim = 4;
+ dl.dtype = pages->dtype;
+ dl.shape = new int64_t[4]{num_pages, num_heads, page_size, head_dim};
+ dl.strides = new int64_t[4]{2 * inner, page_size * head_dim, head_dim, 1};
+ dl.byte_offset = 0;
+ return tvm::ffi::Tensor::FromDLPack(managed, /*require_alignment=*/0,
+ /*require_contiguous=*/false);
+}
+
+/*!
+ * \brief Return a strided, zero-copy view selecting the `[start,
start+length)`
+ * slice along the LAST dimension of \p t, preserving all other strides and
+ * folding the slice offset into the data pointer (so `byte_offset == 0`).
+ *
+ * Used to split MLA tensors that store two head components concatenated along
+ * the last dim: the query into `q_nope`/`q_pe` and the paged cache into
+ * `ckv_cache`/`kpe_cache`. FlashInfer reads tensor strides and ignores
+ * `byte_offset`, so a strided slice avoids a copy.
+ */
+inline ffi::Tensor SliceLastDimView(const Tensor& t, int64_t start, int64_t
length) {
+ int ndim = t->ndim;
Review Comment:
Add a defensive check to ensure that `ndim` is greater than 0. If `ndim` is
0, `in_strides[ndim - 1]` will access `in_strides[-1]`, resulting in an
out-of-bounds memory access.
```suggestion
inline ffi::Tensor SliceLastDimView(const Tensor& t, int64_t start, int64_t
length) {
int ndim = t->ndim;
TVM_FFI_ICHECK_GT(ndim, 0) << "ndim must be greater than 0";
```
##########
src/runtime/vm/attn_backend.h:
##########
@@ -49,6 +49,134 @@ enum class AttnBackendKind : int {
kFlashInfer = 1,
};
+/*!
+ * \brief Return a zero-copy alias of \p t whose `byte_offset` is folded into
the
+ * data pointer, so the resulting tensor has `byte_offset == 0`.
+ *
+ * FlashInfer 0.6.3 kernels read tensors from `data` directly and do NOT honor
+ * the DLPack `byte_offset` field. mlc's auxiliary index tensors (qo_indptr,
+ * kv_indptr, page_indptr, page_indices, length_info, ...) are views packed
into
+ * a shared workspace and therefore carry a non-zero `byte_offset`. Passing
them
+ * as-is makes FlashInfer read the wrong addresses; this helper rebases them.
+ */
+inline ffi::Tensor ZeroByteOffsetView(const Tensor& t) {
+ if (t->byte_offset == 0) return t;
+ auto* holder = new Tensor(t); // keep the underlying storage alive
+ auto* managed = new DLManagedTensor();
+ managed->manager_ctx = holder;
+ managed->deleter = [](DLManagedTensor* self) {
+ delete[] self->dl_tensor.shape;
+ delete[] self->dl_tensor.strides;
+ delete static_cast<Tensor*>(self->manager_ctx);
+ delete self;
+ };
+ DLTensor& dl = managed->dl_tensor;
+ dl.data = static_cast<void*>(static_cast<char*>(t->data) + t->byte_offset);
+ dl.device = t->device;
+ dl.ndim = t->ndim;
+ dl.dtype = t->dtype;
+ dl.shape = new int64_t[t->ndim];
+ dl.strides = nullptr;
+ for (int i = 0; i < t->ndim; ++i) dl.shape[i] = t->shape[i];
+ if (t->strides != nullptr) {
+ dl.strides = new int64_t[t->ndim];
+ for (int i = 0; i < t->ndim; ++i) dl.strides[i] = t->strides[i];
+ }
+ dl.byte_offset = 0;
+ return tvm::ffi::Tensor::FromDLPack(managed, /*require_alignment=*/0,
+ /*require_contiguous=*/false);
+}
+
+/*!
+ * \brief Build a strided, zero-copy view selecting the key (which=0) or value
+ * (which=1) sub-tensor from a combined paged KV tensor of shape
+ * (num_pages, 2, num_heads, page_size, head_dim), yielding a
+ * (num_pages, num_heads, page_size, head_dim) tensor that shares storage with
+ * `pages`. FlashInfer 0.6.3 takes separate key/value paged caches and reads
the
+ * tensor strides, so a strided view avoids an explicit split/copy.
+ */
+inline ffi::Tensor PagedKVCacheView(const Tensor& pages, int64_t which) {
+ TVM_FFI_ICHECK_EQ(pages->ndim, 5);
+ TVM_FFI_ICHECK_EQ(pages->shape[1], 2);
Review Comment:
Add a defensive check to ensure that `which` is either `0` or `1`. If
`which` is out of bounds, it will calculate an incorrect pointer offset (`which
* inner * elem_bytes`), leading to undefined behavior or potential segmentation
faults.
```suggestion
inline ffi::Tensor PagedKVCacheView(const Tensor& pages, int64_t which) {
TVM_FFI_ICHECK(which == 0 || which == 1) << "Invalid which: " << which <<
". Must be 0 or 1.";
TVM_FFI_ICHECK_EQ(pages->ndim, 5);
TVM_FFI_ICHECK_EQ(pages->shape[1], 2);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
