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.