[FlyDSL] [DSv4] Bound the FP4 MQA-logits store with a window-sized V# instead of a compare

[valarLip]

What

The FP4 MQA-logits scorer guarded its per-token store in software:

is_writer = lane_div_16 < 1
in_window = (out_token >= local_start) & (out_token < local_end)
if is_writer & in_window:
    fx.ptr_store(thread_sum, out_base + out_token)

That lowers to 2 v_cmp + 2 s_and_b64 + s_or_b64 + s_and_saveexec_b64 + s_cbranch_execz per n-tile — 28 instructions per chunk, 11.6% of the
steady-state loop body, second only to the reduction itself.

This sizes a buffer descriptor to [local_start, local_end) instead, so
num_records is the window test, in hardware.

A token below the window underflows to a large unsigned offset and is
dropped by the same bound, so one comparison covers both ends.

Two traps, both hit while writing this:

• Lanes 16..63 hold redundant copies of the butterfly result and must not
  write. Parking them past num_records needs a constant offset, not a
  multiple of the window length: token_base - local_start is negative for a
  token below a non-zero window start, and adding the length to that lands
  back inside the window, racing the lane that owns the column. Caught by
  test_flydsl_pa_mqa_logits_fp4_prefill.py's one non-zero-local_start case
  — it went non-deterministic (cos 0.79/0.94/0.99 across runs) where the other
  eight stayed at 1.000000.
• The window bounds must go through readfirstlane. They are wave-uniform but
  arrive from cta_info via a buffer load, i.e. in VGPRs, and a divergent
  descriptor makes the compiler wrap the store in a waterfall loop — 252
  instructions, worse than the guard it replaces.

Both the descriptor and the lane offset are chunk-invariant, so both are built
once outside the loop.

Same fold applied to the rectangular decode twin (pa_mqa_logits_fp4.py).

Numbers

Steady-state loop body 242 → 203 instructions. MI355X (gfx950), 3
interleaved reps, min over 40 iterations:

shape	before	after
rows=8192, W~2064	124.4 us	119.0 us	-4.4%
rows=1024, W~32768	241.0 us	225.4 us	-6.5%

Isolated: only the store path was reverted for the "before" column, so this
is the fold alone and not a mix of unrelated edits.

Correctness

Output is unchanged — cells outside the window still hold the caller's -inf
pre-fill, which test_flydsl_pa_mqa_logits_fp4.py and
test_flydsl_pa_mqa_logits_fp4_prefill.py already assert on every case
(oob_neginf=True, cos_exact=1.000000 / cos_sim=1.0000 on all 13 —
including the non-zero-local_start window set that caught the lane-offset
bug above).

One edge the existing suite does not cover, verified locally but not added
here: a zero-length window (num_records == 0), which only exists as a
concept once the bound is in hardware. Checked with an all-empty-window batch
and with a single live row among empty ones; nothing is written in either.

Test plan

• op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py — 8 cases pass
• op_tests/test_flydsl_pa_mqa_logits_fp4.py — 5 shapes, cos_sim=1.0000
• ISA inspected: buffer_store_dword ... offen, zero v_cmp /
  s_and_saveexec / s_cbranch_execz left in the loop body, no waterfall
  back-edge
• end-to-end accuracy (GSM8K) — not run; the change is bit-identical on
  the in-window cells so I would expect no movement

Note for reviewers

#5282 touches the same file. It rewrites the schedule side
(compute_prefill_schedule, _RowPlan, _row_plan,
_prefill_{row_plan,cta_info}_kernel); this touches only the scoring
kernel's store path. Different regions, so they should merge, but the two
should probably land in a known order.

🤖 Generated with Claude Code

[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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
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 5285 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.