Skip to content

Remove if self.attention != "dot_product": restriction for mla_qk_head_chunk_size in types.py - #4666

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_955775132
Open

Remove if self.attention != "dot_product": restriction for mla_qk_head_chunk_size in types.py#4666
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_955775132

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Remove if self.attention != "dot_product": restriction for mla_qk_head_chunk_size in types.py

Allow mla_qk_head_chunk_size to be used with attention types other than dot_product attention.

@codecov

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@copybara-service copybara-service Bot changed the title FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/AI-Hypercomputer/maxtext/pull/4513 from Shuwen-Fang:sort_one_sparsecore cccf4ecd7360fbb63557537ccc731eb4f5d0b102 Remove if self.attention != "dot_product": restriction for mla_qk_head_chunk_size in types.py Jul 29, 2026
@copybara-service
copybara-service Bot force-pushed the test_955775132 branch 4 times, most recently from b06695f to 1ef1545 Compare July 30, 2026 01:11
@copybara-service
copybara-service Bot force-pushed the test_955775132 branch 3 times, most recently from 666c5f2 to b0143a4 Compare July 30, 2026 02:27
…head_chunk_size` in `types.py`

Allow `mla_qk_head_chunk_size` to be used with attention types other than dot_product attention.

PiperOrigin-RevId: 955775132
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants