[Kernel][FA] Fix ds_read_tr waitcnt hazard and stagger SCC clobber (1/8)

[xinyazhang]

Part 1 of 8 of the gfx950 attention port from AOTriton 0.14.2b production kernel; base main.

Summary

• kernels/attention/flash_attn_utils.py: emit the LDS transpose reads
  (_ds_read_tr16_b64_imm, _ds_read_tr8_b64_imm) as ROCDL ops instead of
  inline asm, so SIInsertWaitcnts sees them and inserts the lgkmcnt wait.
• Same file: tag the dual-wave V reads with the LDS alias scopes the K reads
  already use, so the backend does not insert a vmcnt(0) drain of the KV
  prefetch.
• Same file: declare the ~{scc} clobber on the stagger barrier asm
  (_stagger_extra_barrier_if_zero).
• kernels/common/utils.py: add the signed helpers ssel, smin, smax,
  sdiv_rd_pow2 and vector_elem_type, used by later PRs in the stack.

Both fixes were carried from FlyDSL 0a9c5906 and e87210fe, which lived on a
side branch and never reached main; they were found during the AOTriton 0.14
integration.

Why this fix

The ds_read_tr helpers wrapped the instruction in inline asm. LLVM cannot
see a DS op inside asm, so SIInsertWaitcnts never emitted an
s_waitcnt lgkmcnt before the uses of the result. Once head_dim > 128 spills
to AGPRs, the allocator places v_accvgpr_write copies ahead of the kernel's
own wait and reads registers the LDS load has not yet written. The result is
non-deterministic NaN. With the ROCDL op (pointer formed as a GEP so the
immediate still folds into the DS offset field) the hazard tracker owns the
wait.

The stagger barrier asm runs s_cmp / s_cbranch_scc0, which destroys SCC,
but its constraint string did not say so. A compare scheduled above the asm
could be kept live across it and read stale, which made the LSE store skip.

Performance

MI355X (gfx950), HIP 6 (amd-smi row 6), idle-checked before every launch; bf16; ROCm 10.1, flydsl 0.4.0, torch 2.14.0+rocm10.1.0. Device kernel time (profiler), median TFLOP/s; FLOPs 4BHSqSk*D, halved for causal. Full data: port950-perf-results.md.

Method: each code base runs in its own process and the launches alternate A,B,B,A,A,B,B,A (4 launches each, the table shows the median); inside a process the variants are interleaved with the order flipped every round; 10 rounds x 20 reps per launch. The ± is the launch-to-launch spread (max-min)/median. Both sides call flydsl_flash_attn_func with its defaults (main's old dual-wave kernel); upstream/main 1941889 vs this PR's tip 51aedab.

Config	main	this PR	change
D=128 bf16 causal, S=8192, B=2 H=32	1078 TF (1019.5 us) ±0.9%	1074 TF (1023.4 us) ±0.5%	-0.4%
D=128 bf16 causal, S=8192, B=4 H=8	1054 TF (521.5 us) ±0.9%	1050 TF (523.8 us) ±0.6%	-0.4%
D=128 bf16 causal, S=8192, B=1 H=16	1061 TF (259.0 us) ±0.9%	1057 TF (260.1 us) ±0.6%	-0.4%
D=128 bf16 non-causal, S=8192, B=2 H=32	1140 TF (1928.5 us) ±0.8%	1140 TF (1928.8 us) ±0.5%	-0.0%
D=128 bf16 non-causal, S=8192, B=4 H=8	1140 TF (964.9 us) ±1.0%	1141 TF (963.9 us) ±0.6%	+0.1%
D=128 bf16 non-causal, S=8192, B=1 H=16	1135 TF (484.3 us) ±1.0%	1133 TF (485.3 us) ±0.6%	-0.2%
D=192 bf16 causal, S=8192 (B=2 H=32, B=4 H=8, B=1 H=16)	n/a: does not build	n/a: does not build	n/a

Result: neutral, every row within -0.4% .. +0.1% of main (launch spread about +-1%). The D=192 row cannot be measured:
main and this PR both fail to build it (the generic kernel asserts block_size % threads_per_row_load in
flash_attn_utils.py), so there is no before/after pair.

Validation

• Verified on the branch tip 51aedab3 (MI355X/gfx950, -m "not large_shape"):
  tests/kernels/test_flash_attn_fwd.py: 1033 passed, 2 failed (the two known
  test_paged_fp8_*side_stream* failures, which also fail on unmodified main).
• black and ruff clean on the touched files; python3 scripts/check_repo.py --base upstream/main passes.
• ISA: 192 ds_read_b64_tr_b16, 0 unwaited uses, and the stagger asm
  declares ~{scc}.

Scope

Shared helpers, also used by main's fp8 attention kernels, which get the
fixes too. No kernel algorithm change.

AI model disclosure

• Implementation: Claude Sonnet 5 (agent), orchestrated by Claude Opus 5.5
  (1M context).
• PR description: Claude Sonnet 5.

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

The critical LDS wait hazard lacks automated regression coverage capable of exercising or detecting it.

1 open finding

• Add regression test for D=192 transpose read synchronization · New

What changed in this PR

Fixes gfx950 flash-attention LDS wait hazards and SCC clobber handling while adding signed arithmetic utilities.

Changes:

• Replaces inline-assembly transpose reads with visible ROCDL operations.
• Adds LDS alias scopes and SCC clobber metadata.
• Adds signed integer and vector-type helpers.

File	Description
kernels/​attention/​flash_attn_utils.py	Fixes LDS hazard tracking, aliasing, and SCC clobber declaration.
kernels/​common/​utils.py	Adds signed arithmetic and vector element-type helpers.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[binding7012]

This pr will cause fp8 dense kernel perf regression, such as
B=2,H=16,S=8192, D=192 has a +60% regression @xinyazhang pls take a look

[binding7012]

like these five helper aren't use in this pr? Will this be needed for subsequent PRs?