Adds Grouped and Batched GEMM kernels with blockscaling matching DeepGEMM API

[aryaman-gupta]

Summary

Adds two new grouped and batched FP8 GEMM kernels with blockscaling that mirror the DeepGEMM API on AMD CDNA GPUs:

FlyDSL kernel	DeepGEMM op
kernels/grouped_gemm_blockscale_contiguous.py → compile_grouped_gemm_blockscale_contiguous	m_grouped_fp8_gemm_nt_contiguous
kernels/grouped_gemm_blockscale_masked.py → compile_masked_grouped_gemm_blockscale_masked	m_grouped_fp8_gemm_nt_masked

These ops are core to MoE inference workloads where the gate/up/down projections of multiple experts are batched into a single grouped GEMM. DeepSeek-V3 is the most prominent example — its expert MLPs use FP8 with per-token activation scaling and per-block weight scaling, exactly the configuration these kernels accept. Adding FlyDSL implementations unblocks running such models on AMD hardware via the same call sites that already use DeepGEMM on NVIDIA.

The Python signatures (tensor shapes, dtypes, scale_a / scale_b layouts including the transposed [scale_k, M] activation-scale layout, grouped_layout semantics with -1 for padding, masked_m semantics for the masked variant, (1, 128) × (128, 128) block-scale granularity) are designed to match the DeepGEMM ops byte-for-byte so call-site code can switch backend by import alone.

Test plan

• Unit / correctness: tests/kernels/test_grouped_gemm_blockscale_contiguous.py and ..._masked.py. Coverage spans 1–8 groups, m_per_group from 100 (unaligned) to 1024, plus DeepSeek-V3 shapes (N=2048 K=7168 and N=7168 K=2304) at both bf16 and f16 outputs.
• 30 / 30 tests pass on MI350 (gfx950) at logits_diff_threshold=1e-3.
• Tests registered in tests/arch_compat.py:CDNA_ONLY_TESTS so non-CDNA CI auto-skips.
• pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] for correct CI bucketing.

Run locally:

PYTHONPATH=./ FLYDSL_RUNTIME_ENABLE_CACHE=0 pytest \
  tests/kernels/test_grouped_gemm_blockscale_contiguous.py \
  tests/kernels/test_grouped_gemm_blockscale_masked.py

DeepGEMM conformance details

A few things were chosen explicitly to match DeepGEMM's behavior so any divergence in numerical results vs the NVIDIA path would surface as a real bug rather than a tolerance / convention mismatch:

• Tolerance. The correctness threshold uses the same calc_diff formula as DeepGEMM (cosine-similarity-style normalized error) at logits_diff_threshold=1e-3 — matching DeepGEMM's own FP8 GEMM tests (tests/test_legacy.py:35, tests/test_fp8_fp4.py:194).
• E8M0 quantization on the host. Scale tensors are rounded with ceil_to_ue8m0 (matching DeepGEMM's deep_gemm/utils/math.py:13). Truncation would shrink the scale and cause FP8 saturation on every block; the ceiling is what keeps x / scale_e8m0 ≤ fp8_max.
• Reference convention. Tests compute the reference from pre-quantization FP32 inputs (a_f32 @ b_f32.T) — same as DeepGEMM's tests. The reference contains zero quantization error, so the diff measured against it is the actual end-to-end FP8 → ground-truth error budget.
• Hardware-vs-software scaling, matching DeepGEMM's CUDA / SM strategy. On MI350 (gfx950) the kernel uses the hardware E8M0 path via mfma_scale_f32_16x16x128_f8f6f4, with E8M0 bytes pre-extracted on the host and loaded as uint8 (analogous to DeepGEMM's pack_ue8m0_to_int for the SM100 path). On MI300 (gfx942) where the MFMA-scale instruction is unavailable, scaling is applied in software.

Notes for reviewers

• Deliberate test-style divergence from test_blockscale_preshuffle_gemm.py. Reference uses pre-quantization FP32 instead of dequant-then-matmul, and tolerance is 1e-3 instead of the repo default 2e-3 — both deliberately chosen to match DeepGEMM's convention as described above.

[coderfeli]

@aryaman-gupta flops and Tb/s data compared with aiter ck moe?

[coderfeli]

use normal if instead of scf?

[aryaman-gupta]

050e26d refactors to use the _if_then helper function pattern used by the moe_blockscale_2stage.py. I kept scf instead of normal if for consistency with existing kernels, several of which use scf if for early-exit, including the moe_blockscale_2stage.py kernel.

[coderfeli]

jse fx.Int32 directly?

[aryaman-gupta]

Done in 40f3cfd and 96946d4

[aryaman-gupta]

@aryaman-gupta flops and Tb/s data compared with aiter ck moe?

@coderfeli here are the comparisons with aiter.ck_moe_stage1 from aiter/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common_blockscale.cuh.

TFLOPS (2·M·N·K / time)

E	M_per_E	N	K	FlyDSL masked	FlyDSL contig	CK ck_moe_stage1	masked / CK	contig / CK
8	64	2048	7168	202.3	196.1	149.8	1.35×	1.31×
8	128	2048	7168	392.5	389.9	300.2	1.31×	1.30×
8	256	2048	7168	776.6	766.7	586.8	1.32×	1.31×
8	512	2048	7168	1359.2	1353.2	830.3	1.64×	1.63×
8	64	7168	2048	434.9	410.6	394.8	1.10×	1.04×
8	128	7168	2048	821.2	829.5	519.9	1.58×	1.60×
8	256	7168	2048	992.6	990.2	627.3	1.58×	1.58×
8	512	7168	2048	1357.5	1349.2	743.2	1.83×	1.82×

Bandwidth (TB/s, matmul-input bytes only: (M·K + E·N·K) / time)

E	M_per_E	N	K	FlyDSL masked	FlyDSL contig	CK ck_moe_stage1
8	64	2048	7168	1.63	1.58	1.21
8	128	2048	7168	1.63	1.62	1.25
8	256	2048	7168	1.71	1.68	1.29
8	512	2048	7168	1.66	1.65	1.01
8	64	7168	2048	3.43	3.24	3.11
8	128	7168	2048	3.27	3.30	2.07
8	256	7168	2048	2.01	2.00	1.27
8	512	7168	2048	1.42	1.41	0.78

---

Configuration: gfx950 (MI350), per-1×128 blockscale FP8 inputs (1×128 for activations, 128×128 for weights), BF16 output, balanced expert assignment with topk=1, FlyDSL tile_m=128, CK block_m=64 (CK's production heuristic for tokens > 32). 100 iters / 10 warmup, max_m_factor=1.0 so all Experts receive the same number of tokens. FlyDSL contiguous numbers are essentially identical to masked at this factor; the contiguous kernel's additional advantage on imbalanced/heavily-padded inputs is not exercised here.

TFLOPS: counts only matmul work (2·M·N·K). ck_moe_stage1 additionally fuses a SiLU+multiply epilogue (~0.01–0.04% of matmul FLOPs but ~1–5% of wall-clock time at these shapes); CK TFLOPS is understated by that fraction.

BW: matmul-input bytes only — output writes excluded