[XTile][RaggedDot] Implement Triton group-GEMM (kRaggedDot)

[mfrancepillois]

Implements end-to-end Triton support for kRaggedDot (grouped GEMM) via the experimental XTile emitter, covering all three ragged-dot modes and wiring them through the GEMM autotuner. Gated behind the xla_gpu_experimental_triton_ragged_dot flag (default: false).

--- Tiling analysis (tiling_space.cc/.h, tile_propagation.cc) ---

Extend the experimental tiling system with ragged-dot support:

• kRaggedNonContracting — G is a kSequential outer loop. Output (M_total, N); G and K are kSequential, M and N are kParallel. RTVars: group_size[g] bounds the M sequential loop per group.

• kRaggedContracting — G is kParallel (output dim 0). Grid = G x K_tiles x N_tiles; M is kSequential (inner accumulation). RTVars: group_size[g] bounds the M loop; start_m[g] is its prefix sum.

• kRaggedBatch — regular batched GEMM tiling (B_total, M, N output).

Tile propagation (PropagateTileToInputForRaggedDotOp) derives LHS/RHS/gs operand tiles from the output tile for all three modes.

--- Scheduling (scheduling.cc) ---

L2-aware traversal-order heuristic for ragged-dot:

• kRaggedNonContracting: if M_avg < N, swap M/N traversal so the smaller LHS activation slab reuses the hot RHS column block longer.
• kRaggedContracting: same heuristic over the K/N pair of the G x K x N grid (treating per-group contracting size vs. N output size).

--- XTile emitter (experimental_fusion_emitter.cc, dot_algorithms.cc) ---

Three kernel variants emitted by EmitRaggedDot:

• kRaggedNonContracting: G sequential outer loop over groups; per-group M masking using group_sizes[g]; loop-carried last_m (prefix sum of group sizes); K-loop accumulates the partial dot.

• kRaggedContracting: Grid = G x K_tiles x N_tiles. Each program computes start_m via an O(G) scalar prefix-sum loop, then iterates over M-tiles within its assigned group, accumulating into the (K, N) output tile.

• kRaggedBatch: regular batched GEMM — output[b,:,:] = LHS[b,:,:] @ RHS[b,:,:]. group_sizes is not used in the computation.

GROUP_SIZE L2 tile reordering (ApplyGroupSizeTileIdRemapping):
When BlockLevelFusionConfig.group_size > 1, remaps the flat program_id
before coordinate extraction so that group_size consecutive M (or G)
programs share the same N (or KN) tile, keeping the RHS block hot in L2.
Supports both NonContracting (2-D M x N grid) and Contracting
(3-D G x K x N grid); skipped for kRaggedBatch and batched variants.

--- Autotuner (triton.cc, triton.h) ---

• Route kRaggedDot fusions through the GEMM autotuner using kTritonGemmFusionKind (enables shared profiling / cache infrastructure).

• GetSupportedConfigsForRaggedDot: dedicated search-space generator.

  • Exhaustive: {16,32,64,128,256}^3 x {2,4,8} warps x {1,2} stages x {1,2,4,8} group_size; pruned by dimension size and register budget.
  • Representative (autotune_level > 0, non-exhaustive): 5 curated configs.
  • kRaggedContracting pruning: skip block_m * num_stages >= q_min to avoid pipeline depth exceeding the minimum group size; skip group_size > G.
  • kRaggedNonContracting: skip group_size > num_m_tiles.

• ApplyConfig: write output_tile and inner_tile into the block-level fusion config for all three modes. kRaggedContracting: Grid = G x K x N (output_tile = [G=1, K, N], inner_tile = [M=block_m]).

• PrepareHloModuleForAutotuning: replace group_sizes operand with a balanced constant (floor(M/G) or ceil(M/G)) so autotuner profiling is reproducible across runs with varying input distributions.

--- GEMM rewriter (gemm_rewriter.cc) ---

• Detect kRaggedDot instructions gated behind xla_gpu_experimental_triton_ ragged_dot flag; create kTritonGemmFusionKind custom-call fusions with block_level_fusion_config holding default tile sizes for each mode.

• Default tile sizes:

  • kRaggedNonContracting: output [M, N], inner [G=1, K].
  • kRaggedContracting: output [G=1, K, N], inner [M].
  • kRaggedBatch: output [B=1, M, N], inner [K].

--- Triton fusion / xtile compiler (fusion.cc, xtile_compiler.cc) ---

• Add kTritonGemmFusionKind to the list of XTile-supported fusion kinds so kRaggedDot fusions are emitted via TileAndEmitXTileModule.

• GetNumberOfBlocks: use trailing-dim alignment when tile_sizes is shorter than the output shape, correctly counting blocks for 3-D kRaggedContracting tiles [G=1, BLOCK_K, BLOCK_N] against a (G, K_out, N_out) output shape.

--- Priority fusion (priority_fusion.cc) ---

• Graceful fallback when EstimateRunTimeForTriton returns an error for kRaggedDot fusions (SymbolicTileAnalysis fails on the non-affine G->M prefix-sum mapping); fall back to the standard cost model instead of crashing.

--- Proto / flags ---

• xla.proto: xla_gpu_experimental_triton_ragged_dot — master gate; routes kRaggedDot through XTile instead of hipBLASLt GroupedGEMM.
• backend_configs.proto: group_size in BlockLevelFusionConfig — L2 tile reordering group size for group-GEMM kernels (0/1 = off).
• autotuning.proto: group_size in TritonConfig — same semantics, used by the autotuner result cache.

--- Indexing analysis (indexing_analysis.cc) ---

• Handle kRaggedDot in HloInstructionIndexing so the compiler can propagate indexing maps through ragged-dot fusions.

--- Tests (gemm_rewriter_triton_ragged_dot_test.cc) ---

End-to-end correctness tests:

• kRaggedNonContracting: balanced/unbalanced groups, constant group_sizes, FP16/BF16/FP8, batched, non-multiple block size, various layouts.
• kRaggedContracting: balanced/unbalanced groups, non-multiple block, LHS transposed, FP16/BF16/FP8, output column-major, S64 group_sizes.
• Autotuner integration: both modes with all-configs exhaustive pass and LargeContracting/LargeNonContracting benchmarks.

[claude]

Review Summary

This PR implements end-to-end Triton support for kRaggedDot (grouped GEMM) via the XTile backend, covering all three ragged-dot modes (NonContracting, Contracting, Batch). The feature is gated behind xla_gpu_experimental_triton_ragged_dot (default: false) and spans tiling analysis, scheduling heuristics, the XTile emitter, autotuner search space, GEMM rewriter, and comprehensive end-to-end tests.

Key findings (see inline comments):

• gpu_profiler.cc: Early return absl::OkStatus() fires for all fusions with 3+ operands, not just those containing kRaggedDot. Currently benign (kFusion/kCustomCall are mutually exclusive opcodes) but fragile for future handlers. Suggestion provided.
• gpu_profiler.cc: Hardcoded buffer_index=2 for group_sizes relies on an implicit contract between inner ragged dot operand ordering and outer fusion operand ordering. Suggest adding an assertion to verify the mapping.
• triton.cc: ~20 lines of dead code computing gs_literal that is immediately discarded with (void)gs_literal. The explanatory comment is valuable; the computation is not.
• experimental_fusion_emitter.cc: Hardcoded F32 accumulator type diverges from the regular EmitDot path which uses GetDotAccumulatorType(). If this is intentional, it should be documented; if not, it should use the same API.
• autotuner_pass.cc: The new ShouldAutotuneInstruction guard changes behavior for all generic fusions, not just ragged-dot. Worth noting in the PR description.

🤖 Generated with Claude Code