Skip to content

Cast tl.program_id(0) to int64 in llama4_rope and qwen2vl_mrope - #1343

Open
akshansh47 wants to merge 1 commit into
linkedin:mainfrom
akshansh47:fix/int64-program-id-llama4-qwen2vl-mrope
Open

Cast tl.program_id(0) to int64 in llama4_rope and qwen2vl_mrope#1343
akshansh47 wants to merge 1 commit into
linkedin:mainfrom
akshansh47:fix/int64-program-id-llama4-qwen2vl-mrope

Conversation

@akshansh47

Copy link
Copy Markdown

Summary

  • Port the #804 int64 tl.program_id(0) cast to two kernels that missed it: llama4_rope.py and qwen2vl_mrope.py.
  • Same int32 overflow hazard as rope.py / rms_norm.py; Ascend llama4_rope already had .to(tl.int64).

Test plan

  • Existing test/transformers/test_llama4_rope.py and test_qwen2vl_mrope.py on GPU CI
  • Diff is two one-line casts matching the #804 pattern

Fixes #1335

Made with Cursor

Same int32 overflow hazard fixed for rope.py / rms_norm.py in linkedin#804.
llama4_rope and qwen2vl_mrope missed the cast; Ascend llama4_rope already
has it. Fixes linkedin#1335.

Co-authored-by: Cursor <[email protected]>
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.

llama4_rope / qwen2vl_mrope kernels missing the int64 program_id cast that #804 added to rope.py/rms_norm.py

1 participant