[FlyDSL] jagged_dense_bmm_broadcast_add (MI300X)

[anhminhnguyenhoang]

Motivation

jagged_dense_bmm_broadcast_add is a grouped GEMM for recommendation-style workloads with variable per-group sequence lengths. For each group b over its packed row slice [s, e):

Out[s:e, :] = Jagged[s:e, :] @ Dense[b] + Bias[b][None, :]

Group row boundaries come from a device-resident seq_offsets prefix-sum array — the host does not know each group's row count at launch, so the group→row mapping is resolved on the GPU. This rules out any stock batched-GEMM kernel.

This PR adds a FlyDSL implementation for gfx942 (MI300X) that delivers ~1.3–1.6× over the generative-recommenders Triton baseline on four headline deployment shapes (B120/B1024 × D256/D512, Mi=7680), plus a tuned winner for the production shape B64 D512 KOUT1024 Mi8192.

Files

Path	Role
aiter/ops/flydsl/kernels/jagged_dense_bmm_gen.py	BF16 grouped GEMM kernel
aiter/ops/flydsl/kernels/jdbba_skew_tile_map.py	Device-side TILE_MAP prep for skew launches
aiter/ops/flydsl/kernels/_buffer_utils.py	Bounded buffer helper for per-group OOB safety
aiter/ops/flydsl/jagged_dense_bmm_dispatch.py	Public dispatch API (jagged_dense_bmm_dispatched)
aiter/ops/flydsl/jagged_dense_bmm_dispatch.json	Arch-keyed autotune winners (gfx942)
op_tests/test_jagged_dense_bmm.py	Perf + correctness vs Triton, skew TILE_MAP, dispatch routing

Technical Details

What it computes

Each group b multiplies its jagged rows by a per-group dense weight matrix and adds a per-group bias. Row boundaries live in seq_offsets on the GPU — the kernel reads them at launch rather than relying on a fixed batch size.

Tensor	Shape	Role
A (jagged)	(L, K) bf16	Packed input rows (L = seq_offsets[-1])
B (dense)	(B·N, K) bf16	One N×K weight panel per group, stacked tall
BIAS	(B·N,) bf16	Broadcast bias per output column
SEQ_OFFSETS	(B+1,) int32	Prefix sum of per-group row counts
C (out)	(L, N) bf16	Packed output

Entry point: jagged_dense_bmm_dispatched(C, A, B, BIAS, SEQ_OFFSETS, n_groups, max_seq_len, uniform_seqlen=...).

Kernel overview

A tiled BF16 GEMM on MI300X using 16×16×16 MFMA. Default tile size is 128×128×64 with 256 threads.

• Jagged input (A) is staged through LDS with double-buffering and bank-conflict swizzling.
• Dense weights (B) are prefetched directly into registers (3-stage pipeline) to hide memory latency.
• Accumulation is FP32; bias is fused in FP32 before writing bf16 output.
• Per-group safety: each block only processes rows belonging to its group. Partial tail tiles and empty groups are handled via bounded buffers and early exit.
• AOT compile cache: first launch per config key is compiled via flyc.compile; results are cached in a bounded LRU (cap 64). Compile failures clean up leaked ir.Context and re-raise (mirrors moe_kernels._run_compiled).

Two launch modes

Mode	When	Idea
Uniform	Every group has the same length	Regular 3D grid over (group, M-tile, N-tile)
Compact (skew)	Variable per-group lengths	Build a TILE_MAP on device listing only occupied tiles; skip empty work

For skewed lengths, a fused prep kernel scans seq_offsets, fills a compact tile list (with sentinels for padding), and the main kernel iterates that list instead of launching empty blocks.

Tuning and dispatch

Per-shape settings (XCD grid remap, thread count, MFMA variant, block_k) are stored in jagged_dense_bmm_dispatch.json under gfx942. Dispatch picks config in order:

1. Explicit call-site overrides
2. Exact JSON match on shape key B{n}D{K}K{N}N{max_seq_len}
3. Heuristic fallback by reduction dimension

Key tuning choices on gfx942:

• BLOCK_K=64 — keeps LDS usage low enough for good occupancy (vs 128, which fills 64 KB and halves occupancy).
• threads=512 — helps D256 uniform shapes by reducing register pressure; D512 and skew stay at 256.
• XCD remap — reorders tile scheduling across chiplets for better L2 reuse; per-shape (xcd_c, xcd_w) winners in JSON. Uniform and compact (skew) paths share one _xcd_remap helper.
• block_k override — wired through dispatch JSON; explicit values are validated before compile (multiple of 32, ≥64 for default MFMA path, LDS fit, K % block_k == 0). Supported values for typical shapes: 64, 128.

Override JSON path: FLYDSL_JAGGED_DENSE_BMM_DISPATCH_JSON. Force arch section: FLYDSL_JAGGED_DENSE_BMM_ARCH.

Review fixes

Addressed reviewer feedback (coderfeli):

• Merged _xcd_remap / _xcd_remap_compact into one shared helper (compile-time vs runtime num_rows).
• Replaced functools.lru_cache on the launcher with flyc.compile + bounded LRU cache.
• Rewrote jdbba_skew_tile_map.py to idiomatic FlyDSL (buffer_ops, loop-carried range) — removed raw MLIR (scf/arith/InsertionPoint).
• Removed unused non-fused tile-map path (build_tile_map_device, scatter launcher).
• Wired block_k fully through dispatch; added fail-fast validation and regression tests.

Test Plan

MI300X (gfx942), ROCm 6.x+. Requires aiter, flydsl ≥0.2.4, and the baseline Triton kernel from generative-recommenders on PYTHONPATH.

# Build (full rebuild on first run or after C++ changes):
AITER_REBUILD=1 python -c "import aiter"

# Full sweep — perf + correctness (uniform/skew) + dispatch routing + regression tests:
HIP_VISIBLE_DEVICES=5 \
PYTORCH_ALLOC_CONF=expandable_segments:True \
PYTHONPATH=<aiter>:<generative-recommenders>:$PYTHONPATH \
    python op_tests/test_jagged_dense_bmm.py

# Headline shapes only (uniform + skew):
HIP_VISIBLE_DEVICES=5 \
PYTORCH_ALLOC_CONF=expandable_segments:True \
PYTHONPATH=<aiter>:<generative-recommenders>:$PYTHONPATH \
    python op_tests/test_jagged_dense_bmm.py \
        -s 120,256,256,7680 120,512,512,7680 1024,256,256,7680 1024,512,512,7680

# Production shape (uniform + skew):
HIP_VISIBLE_DEVICES=5 \
PYTORCH_ALLOC_CONF=expandable_segments:True \
PYTHONPATH=<aiter>:<generative-recommenders>:$PYTHONPATH \
    python op_tests/test_jagged_dense_bmm.py \
        -s 64,512,1024,8192 -r uniform skew

The op test verifies:

1. All headline shapes pass correctness (checkAllclose vs torch reference, err=0).
2. Skew regime — empty groups (M_b=0), partial tiles, max_seq_len >> mean — via compact TILE_MAP.
3. In-table shapes resolve to the JSON winner; off-table shapes use the D-bucketed heuristic (test_jdbba_dispatch).
4. block_k=128 correctness and fail-fast ValueError for invalid block_k (test_jdbba_block_k).
5. Skew varying L across different seq_offsets (test_jdbba_skew_varying_L).
6. flyc.compile failure re-raises without poisoning the compile cache (test_jdbba_compile_reraise).

Test Results

MI300X gfx942, ROCm 7.2, FlyDSL 0.2.4, Triton 3.7.0. run_perftest timing, checkAllclose (err=0) all shapes/regimes.
Baseline: generative_recommenders.ops.triton.triton_jagged.triton_jagged_dense_bmm_add_fwd.

Uniform regime (every group M_i = Mi)

B	D	KOUT	Mi	FlyDSL (ms)	Triton (ms)	Speedup
120	256	256	7680	0.494	0.664	1.34×
120	512	512	7680	1.450	1.977	1.36×
1024	256	256	7680	4.262	6.004	1.41×
1024	512	512	7680	12.252	17.072	1.39×
64	512	1024	8192	1.533	2.226	1.45×

Production uniform (B64, M_total=524288): JSON winner xcd_c=240, xcd_w=4.

Skew regime (~20% empty groups, one full-envelope, one near-full)

B	D	KOUT	Mi	FlyDSL (ms)	Triton (ms)	Speedup
120	256	256	7680	0.090	0.140	1.56×
120	512	512	7680	0.261	0.389	1.49×
1024	256	256	7680	0.693	1.080	1.56×
1024	512	512	7680	1.992	3.178	1.60×
64	512	1024	8192	0.314	0.475	1.51×

Production skew (B64): compact TILE_MAP with _skew_compact_xcd remap (xcd_c=32, xcd_w=8, D≥512).

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 4136 --add-label <label>