Skip to content

fix XFAIL#49

Open
toudefu wants to merge 2 commits into
MetaX-MACA:devfrom
toudefu:dev
Open

fix XFAIL#49
toudefu wants to merge 2 commits into
MetaX-MACA:devfrom
toudefu:dev

Conversation

@toudefu

@toudefu toudefu commented Jul 21, 2026

Copy link
Copy Markdown

fix: tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +685 to +690
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;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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;

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant