Repository navigation
[FlyDSL] fp8_mqa_logits: drop the host row-pad copies - #6001
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
Tail masking changes GPU memory addressing in two optimized kernels and warrants final human and hardware-backed validation.
Review effort: Balanced
Findings: 1
Open (1)
What changed in this PR
Removes host-side row padding by safely masking short tail tiles in both FlyDSL kernel variants.
Changes:
- Clamps tail-row loads and suppresses invalid-row writes.
- Allocates and launches using the real sequence length.
- Updates short-tail test coverage and documentation.
| File | Description |
|---|---|
fp8_mqa_logits.py |
Adds tail-row masking to both kernel builders. |
fp8_mqa_logits_kernels.py |
Removes host padding and passes real dimensions. |
test_flydsl_fp8_mqa_logits.py |
Updates short-tail test context. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
ChuanLi1101
left a comment
ChuanLi1101
left a comment
There was a problem hiding this comment.
Removing the host padding is safe. The kernel already launches one block for each partial tile of rows. For the extra slots in the last tile, the loads stay on the last real row and the KV window is empty, so nothing is stored. The -inf fill is turned off for those slots too, so it does not wipe the last real row. The trace matches this: the host copies and zero-fills disappear, and the kernel time stays the same.
I am not approving the PR yet. The tests had not started, so I am approving the Aiter Test, Checks, and OPUS workflow runs. Adding a CI label will not start them. Once #5921 merges, the split count still has to be computed from the real row count, rounded up, not from the old padded length.
vpietila-amd
left a comment
vpietila-amd
left a comment
There was a problem hiding this comment.
Looks good to me. Thanks for fixing this.
…950-epilogue Brings in #6001: the launcher no longer zero-pads Q, weights, cu_starts and cu_ends; both kernel builders mask the short last row tile. Conflicts in fp8_mqa_logits_kernels.py: kept this branch's occupancy-aware / aggressive split logic and passed it the unpadded seq_len. Co-Authored-By: Claude <noreply@anthropic.com>
The kernel already launches ceil(seq_len / rows_per_block) blocks. Padding Q, weights, and the window up to that multiple launched four zero-fills and four cats on every unaligned call, and the extra rows were never read. Mask the short last tile and pass the real seq_len. Signed-off-by: Frida Andersson <fanderss@amd.com>
An empty window already keeps a past-the-end slot from being stored, so the extra writer predicate and the epilogue pointer clamp were not doing any work. Signed-off-by: Frida Andersson <fanderss@amd.com>
Black 26.5.1 wants the short _make_out_row_t call on one line. That is the pre-check that blocked the test gate. Signed-off-by: Frida Andersson <fanderss@amd.com>
The auto-variant pin still said (257, 8192) pads. The launcher no longer pads, so the note should say the case leaves a short tail tile. Signed-off-by: Frida Andersson <fanderss@amd.com>
e7d0738
e6d4da4 to
e7d0738
Compare
The rebase event still had the old changed-file count, so the tag job compared 3 files with 744. Signed-off-by: Frida Andersson <fanderss@amd.com>
|
Rebased onto main now that #5921 is merged. The split count uses the real Could you LGTM again? @ChuanLi1101 @vpietila-amd @zhiding512 |
…tical to main The merge with main reintroduced host-side padding of Q/weights/cu_starts/ cu_ends to a multiple of the block's rows for every variant outside the gfx942 LDS set. main dropped that in ROCm#6001 (every builder masks a short last block in-kernel), so on gfx950 an unaligned seq_len went from one launch to nine and regressed (reported on MI355X). Pass the real seq_len to every variant again, as main does; only the split policy stays variant-specific (_LDS942_SPLIT_POLICY for the gfx942 LDS variants). Also hard-gate the PR to gfx942: revert the shared direct-load builder (_build_kernel_mfma_r_w, used by gfx950's direct-load variants) to main's version. The empty-row union-window change only existed to absorb padded rows, and the hoisted _fn_to_fnuz_i64 now lives in the gfx942 LDS module, which is its only other user. fp8_mqa_logits.py is now byte-identical to main. Add a launch-free launcher-contract check to the op test (real seq_len, no host cat/new_zeros) that runs on every arch. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
Summary
Remove the host-side row padding from
flydsl_fp8_mqa_logits.For an unaligned
seq_len, the launcher used to zero-fill and concatenate Q, weights,cu_starts, andcu_ends. The kernel already uses a ceil-div row grid, so it can handle the short last tile directly.Both kernel builders now clamp tail loads to the last valid row and give those slots an empty window. The launcher passes the real
seq_len;_auto_variantis unchanged.Trace
DeepSeek-V4.1-Flash, TP4, MI325X, rank 0. Two 262144 → 1024 requests, concurrency
4, DSpark-5, FlyDSL QuickReduce.
Before (
M=16370,mfma_r4_w4):After:
The rank-0 capture had 971 padded launches. This removes 3884 cats (25.4 ms) and 3884 fills (17.5 ms). There are no pad cats or fills in the new trace.
The kernel time is flat:
mfma_r4_w4averages 1.124 ms before and 1.129 ms after. The alignedM=16380step did not pad and is also flat (478.1 ms before, 473.6 ms after, median of 16).Testing
257×8192and16370×32768, including empty and past-end windows