feat(ep): optimize warp id allocation under dispatch ll kernel

[kawhil-amd]

Summary

• Rotate header warp selection across SIMD units: instead of always assigning inGroupWarpId == 0 as the header warp, rotate the selection so header
  work (atomic allocation, shared-memory signaling) spreads evenly across all 4 SIMD units, avoiding SIMD0 hotspotting.
• Vectorize shared-memory initialization: replace per-group serial clears (inGroupWarpId == 0 && laneId == 0 loop) with a single warp-wide clear
  (warpId == 0 && laneId < kMaxWarpGroups), cutting init latency.
• Hoist tokenIndices lookup before the duplicate-dest branch: compute tokExpert once and reuse it in the duplicate-destination check, eliminating
  a redundant global load per iteration.
• Move dispatchOut prefetch into the header/normal-warp branches: the remote address fetch is now co-located with its consumer — header warp fetches
  after allocation, normal warps prefetch while spin-waiting, improving occupancy by overlapping fetch with useful work.
• Increase WarpCopy unroll factor for bf16: bump from 2 to 4 for hip_bfloat16, doubling the bytes moved per unrolled iteration to better utilize
  memory bandwidth at the same instruction count.

Dispatch Latency (MI355X)

block_num	topk	max_token	Baseline(us)	Current(us)	Speedup
32	8	32	19.16	18.56	+3.1%
64	8	64	24.04	23.44	+2.6%

[jhchouuu]

Reviewed against the full EpDispatchIntraNodeLLKernel_body on main. Logic looks correct for the current kWarpsPerGroup == 2 config. Nice, well-targeted micro-opt. A few notes below.

Verified correct

• Header rotation: headerWarpIdx ∈ [0, kWarpsPerGroup), so each group still has exactly one header (producer) and the rest are consumers. groupData/groupCounters are keyed by warpGroupIdInBlock, independent of which inGroupWarpId is the header, so the producer/consumer protocol is preserved. For k=2 the else-branch yields header SIMDs {0,2,1,3} as the comment claims.
• Vectorized shared-mem clear: warpId == 0 && laneId < kMaxWarpGroups clears all 8 slots in one warp; __syncthreads() still guards visibility. Clearing unused slots is harmless.
• tokExpert hoist: for the lanes that matter (laneId < i % numExpertPerToken < numExpertPerToken), laneId % numExpertPerToken == laneId, so the value matches the original tokenIndices[base + laneId]; the % only keeps the other lanes in-bounds. (Note: this hoists/broadcasts the load rather than removing one — load count is unchanged, it's a latency-hiding move.)
• dispatchOut prefetch sink: the skip path continues before use, and both header/normal branches assign dispatchOut before the WarpCopy, so the initial nullptr is never dereferenced.

Suggestion (robustness)

The skip-marker writer changed from the header warp to !isHeaderWarp:

// "only warp 0 writes the skip marker"  <-- comment is now stale
if (!isHeaderWarp && laneId == 0) {
  args.dispDestTokIdMap[i] = FlatTokenIndex(config, config.worldSize, 0);
}

For kWarpsPerGroup == 2 this selects exactly one (non-header) warp — correct. But if kWarpsPerGroup is ever raised (the kWarpsPerGroup >= kNumSimds branch is clearly anticipating that), !isHeaderWarp becomes true for multiple warps, all writing the same constant to the same dispDestTokIdMap[i] — idempotent so not a correctness bug, but redundant and the "only warp 0 writes" comment no longer holds. Consider either pinning a single fixed writer (e.g. inGroupWarpId == 0 && laneId == 0, decoupled from the header role) or adding static_assert(kWarpsPerGroup == 2) to nail down the assumption.

Nits

• The header→SIMD mapping (warpId % 4) is a soft scheduling assumption, not architecturally guaranteed; if it doesn't hold the spread is just less even — no correctness impact. Worth a one-line "best-effort" note in the comment.
• A few case-only churn changes add diff noise without behavior change: assert text "Total recv..."→"total recv...", comment Normal Warps→normal Warps. Consider reverting. (The Each warp-group (4 warps) → Each warp-group fix is good — the old "4 warps" was wrong for k=2.)