fix XFAIL#49
Conversation
There was a problem hiding this comment.
Code Review
This pull request replaces alignas with __attribute__((aligned(...))) in the MACA codegen, enables previously xfailed paged attention tests, and updates CopyTreeAttnMaskOnDepthAsync to return a 2D view. Feedback suggests optimizing CopyTreeAttnMaskOnDepthAsync by avoiding the creation of two separate views to prevent extra heap allocations, and adding a defensive check to ensure the data size is even.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| Tensor view_1d = | ||
| tree_attn_mask_device_[depth].CreateView({static_cast<int64_t>(data->size())}, dtype_aux_); | ||
| CopyVecDataToArray(view, data->data()); | ||
| return view; | ||
| Tensor view_2d = tree_attn_mask_device_[depth].CreateView( | ||
| {static_cast<int64_t>(data->size()) / 2, 2}, dtype_aux_); | ||
| CopyVecDataToArray(view_1d, data->data()); | ||
| return view_2d; |
There was a problem hiding this comment.
Instead of creating two separate views (view_1d and view_2d), we can create only view_2d and pass the expected 1D shape directly to CopyVecDataToArray. This avoids an extra heap allocation for the temporary view_1d tensor and its shape on every forward pass of tree attention. Additionally, we should defensively check that data->size() is even before dividing by 2.
| Tensor view_1d = | |
| tree_attn_mask_device_[depth].CreateView({static_cast<int64_t>(data->size())}, dtype_aux_); | |
| CopyVecDataToArray(view, data->data()); | |
| return view; | |
| Tensor view_2d = tree_attn_mask_device_[depth].CreateView( | |
| {static_cast<int64_t>(data->size()) / 2, 2}, dtype_aux_); | |
| CopyVecDataToArray(view_1d, data->data()); | |
| return view_2d; | |
| TVM_FFI_ICHECK_EQ(data->size() % 2, 0); | |
| Tensor view_2d = tree_attn_mask_device_[depth].CreateView( | |
| {static_cast<int64_t>(data->size()) / 2, 2}, dtype_aux_); | |
| CopyVecDataToArray(view_2d, data->data(), ffi::Shape({static_cast<int64_t>(data->size())})); | |
| return view_2d; |
fix: tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py