Chefang mha global load

[fangche123]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4357 --add-label <label>

[Copilot]

Pull request overview

Adds support and validation utilities for very-large-sequence / very-large-KV FMHA v3 forward execution, including a host-side workaround for the 65,535 limit on grid.z in gfx950 varlen group launches.

Changes:

• Add two forward-only correctness scripts targeting extreme sizes: (>4 GiB per-batch K/V) and (very large seqlen with last-row checking).
• Extend fmha_fwd_v3_args with s_q_tile_base and initialize it on the host.
• Split gfx950 varlen group-mode launches into multiple grid.z chunks when Q-tiles exceed 65,535, passing s_q_tile_base to the kernel.

Reviewed changes

Copilot reviewed 4 out of 8 changed files in this pull request and generated 2 comments.

File	Description
op_tests/test_mha_varlen_large_kv.py	New script to validate large-seqlen varlen/group FMHA v3 forward by checking only the last query rows.
op_tests/test_mha_fwd_large_kv.py	New script to validate >4 GiB per-batch K/V addressing correctness in FMHA v3 forward.
csrc/include/mha_fwd.h	Adds s_q_tile_base to the packed v3 forward kernel-arg struct.
csrc/cpp_itfs/mha_fwd.cu	Initializes s_q_tile_base and chunks gfx950 varlen group launches across multiple grid.z launches.

Comments suppressed due to low confidence (1)

op_tests/test_mha_fwd_large_kv.py:115

• The 4GiB-path warning currently triggers only when both K and V are <= 4GiB, but the docstring says both must exceed 4GiB to exercise the boundary. If either one stays <=4GiB, the test won't validate that tensor's >4GiB addressing.

    if k_bytes <= _4GiB and v_bytes <= _4GiB:
        print("WARNING: neither K nor V exceeds 4 GiB; this will NOT exercise the "
              ">4GiB path. Increase -k / -n.")

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.