diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 98d78bc50..0fb36a0b7 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -266,6 +266,10 @@ def _test_grouped_linear_accuracy( split_size = 1 if fp8: split_size = get_align_size_for_quantization(recipe) + elif os.environ.get("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1": + # The FlyDSL grouped GEMM tiles every expert with BLOCK_M=256, so even the + # unquantized path needs 256-aligned per-expert token counts. + split_size = 256 m = config.max_seqlen_q // split_size dist = torch.sort(torch.randint(0, m, (num_gemms - 2,))).values.tolist() dist.append(dist[-1]) # Manually add a zero @@ -598,9 +602,9 @@ def test_grouped_linear_accuracy_cutlass( ) -# ROCm grouped GEMM backends (HipKittens / CK). Kept separate from the CUTLASS test above so the -# upstream one stays untouched. The backend is picked by env var; use_cutlass only relaxes the -# tolerances, since a non-reference GEMM will not be bit-exact against the sequential Linear. +# ROCm grouped GEMM backends (HipKittens / CK / FlyDSL). Kept separate from the CUTLASS test above +# so the upstream one stays untouched. The backend is picked by env var; use_cutlass only relaxes +# the tolerances, since a non-reference GEMM will not be bit-exact against the sequential Linear. @pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm-only grouped GEMM backends") @pytest.mark.parametrize("dtype", param_types, ids=str) @pytest.mark.parametrize("num_gemms", [3, 6]) @@ -610,7 +614,7 @@ def test_grouped_linear_accuracy_cutlass( @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) @pytest.mark.parametrize("delay_wgrad_compute", all_boolean) -@pytest.mark.parametrize("grouped_gemm_backend", ["hipkittens", "ck"], ids=str) +@pytest.mark.parametrize("grouped_gemm_backend", ["hipkittens", "ck", "flydsl"], ids=str) def test_grouped_linear_accuracy_rocm_backends( dtype, num_gemms, @@ -638,13 +642,32 @@ def test_grouped_linear_accuracy_rocm_backends( # Triton backends in test_grouped_linear_accuracy instead. pytest.skip("HipKittens/CK grouped GEMM backends do not support FP8 block scaling.") + if grouped_gemm_backend == "flydsl": + # FlyDSL grouped GEMM is gfx950-only and covers MXFP8, tensor-wise FP8 + # (current/delayed scaling), and unquantized bf16/fp16. Skip configs its gate + # rejects, which would silently fall back to the default backend and give + # false "flydsl verified" coverage. + if get_device_compute_capability() != (9, 5): + pytest.skip("FlyDSL grouped GEMM is only supported on gfx950.") + if recipe is None and dtype not in (torch.bfloat16, torch.float16): + pytest.skip("FlyDSL unquantized grouped GEMM supports only bf16/fp16.") + if recipe is not None and not ( + recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed() + ): + pytest.skip("FlyDSL grouped GEMM supports MXFP8 and tensor-wise FP8 recipes only.") + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.delenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_CK_GROUPED_GEMM", raising=False) + monkeypatch.delenv("NVTE_USE_GROUPED_GEMM_FLYDSL", raising=False) if grouped_gemm_backend == "hipkittens": monkeypatch.setenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", "1") - else: + elif grouped_gemm_backend == "ck": monkeypatch.setenv("NVTE_USE_CK_GROUPED_GEMM", "1") + else: + # FlyDSL has its own opt-in; keep the cutlass/HK/CK grouped GEMM path disabled. + monkeypatch.delenv("NVTE_USE_CUTLASS_GROUPED_GEMM", raising=False) + monkeypatch.setenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "1") # For the configs CK grouped GEMM is guaranteed to handle, # assert it does not silently fall back to multi-stream hipBLASLt. @@ -656,6 +679,32 @@ def test_grouped_linear_accuracy_rocm_backends( if expect_ck_no_fallback: monkeypatch.setenv("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", "1") + # FlyDSL falls back to the default grouped GEMM silently on any config it rejects + # past the gate skipped above, and wgrad is not attempted at all with fused or + # delayed wgrad, so a missing fallback warning does not show that FlyDSL ran. + # Record each FlyDSL dispatch and assert on the phases that actually ran; the + # fallback warning stays on so a failure shows the reason. + flydsl_dispatches = [] + if grouped_gemm_backend == "flydsl": + from transformer_engine.pytorch.flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + gemm_wrappers, + ) + + monkeypatch.setenv("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "1") + grouped_gemm_flydsl = gemm_wrappers._grouped_gemm_flydsl + phases = {"TN": "fwd", "NN": "dgrad", "NT": "wgrad"} + + def recording_grouped_gemm_flydsl(*args, layout, **kwargs): + try: + grouped_gemm_flydsl(*args, layout=layout, **kwargs) + except FlyDSLUnsupportedError: + flydsl_dispatches.append((phases[layout], False)) + raise + flydsl_dispatches.append((phases[layout], True)) + + monkeypatch.setattr(gemm_wrappers, "_grouped_gemm_flydsl", recording_grouped_gemm_flydsl) + test_grouped_linear_accuracy( dtype, num_gemms, @@ -671,6 +720,20 @@ def test_grouped_linear_accuracy_rocm_backends( use_cutlass=True, ) + if grouped_gemm_backend == "flydsl": + expected = {"fwd", "dgrad"} + if not fuse_wgrad_accumulation and not delay_wgrad_compute: + expected.add("wgrad") + ran = {phase for phase, ok in flydsl_dispatches if ok} + fell_back = sorted({phase for phase, ok in flydsl_dispatches if not ok}) + assert not fell_back and expected <= ran, ( + f"FlyDSL grouped GEMM did not run every expected phase: expected {sorted(expected)}, " + f"ran {sorted(ran)}, fell back in {fell_back} (dtype={dtype}, " + f"recipe={recipe_id(recipe)}, fp8_model_params={fp8_model_params}, " + f"fuse_wgrad={fuse_wgrad_accumulation}, delay_wgrad={delay_wgrad_compute}). " + "See the [FLYDSL WARNING] for the fallback reason." + ) + if expect_ck_no_fallback: # NVTE_WARN writes to std::cerr; capfd captures file-descriptor-level # output, including C/C++ stderr. @@ -682,6 +745,120 @@ def test_grouped_linear_accuracy_rocm_backends( ) +@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm-only grouped GEMM backend") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize( + "recipe", [recipe.Float8CurrentScaling(), recipe.DelayedScaling()], ids=recipe_id +) +@pytest.mark.parametrize("fp8_model_params", all_boolean) +def test_grouped_linear_flydsl_fp8_scaling(recipe, fp8_model_params, monkeypatch): + """FlyDSL tensor-wise FP8 grouped GEMM vs. sequential Linear over several steps. + + Delayed scaling quantizes its first step with unit scales (there is no amax history + yet), so the single-step accuracy test cannot catch a per-expert scale bug for it. + Run a few steps with uneven experts and require the compared step to hand FlyDSL + non-unit scales that differ across experts, for every phase. + """ + if get_device_compute_capability() != (9, 5): + pytest.skip("FlyDSL grouped GEMM is only supported on gfx950.") + pytest.importorskip("flydsl", reason="FlyDSL package is not installed.") + from transformer_engine.pytorch.flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + gemm_wrappers, + ) + + num_gemms, num_steps, dtype = 3, 3, torch.bfloat16 + config = model_configs["126m"] + monkeypatch.setenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "1") + monkeypatch.setenv("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "1") + + # Record each FlyDSL grouped GEMM: its phase, whether it ran, and the operands' + # per-tensor inverse scales. + phases = {"TN": "fwd", "NN": "dgrad", "NT": "wgrad"} + dispatches = [] + grouped_gemm_flydsl = gemm_wrappers._grouped_gemm_flydsl + + def recording_grouped_gemm_flydsl(A, B, out, *, layout, **kwargs): + scales = {float(t._scale_inv) for t in list(A) + list(B)} + try: + grouped_gemm_flydsl(A, B, out, layout=layout, **kwargs) + except FlyDSLUnsupportedError: + dispatches.append((phases[layout], False, scales)) + raise + dispatches.append((phases[layout], True, scales)) + + monkeypatch.setattr(gemm_wrappers, "_grouped_gemm_flydsl", recording_grouped_gemm_flydsl) + + with quantized_model_init(enabled=fp8_model_params, recipe=recipe): + grouped_linear = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ) + sequential_linear = torch.nn.ModuleList( + [ + Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ) + for _ in range(num_gemms) + ] + ) + with torch.no_grad(): + for i in range(num_gemms): + sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) + + # Uneven, FlyDSL-aligned experts so each carries its own scales. + align = get_align_size_for_quantization(recipe) + tokens = config.max_seqlen_q + m_splits = [tokens // 4, tokens // 2, tokens // 4] + assert sum(m_splits) == tokens and all(m % align == 0 for m in m_splits) + reset_rng_states() + inputs = [ + torch.randn(tokens, config.hidden_size, dtype=dtype, device="cuda") + for _ in range(num_steps) + ] + + def run(block): + FP8GlobalStateManager.reset() + for inp in inputs: + dispatches.clear() + for p in block.parameters(): + p.grad = None + inp = inp.clone().requires_grad_() + with autocast(enabled=True, recipe=recipe): + if isinstance(block, GroupedLinear): + out = block(inp, m_splits) + else: + out = torch.cat([block[i](x) for i, x in enumerate(torch.split(inp, m_splits))]) + out.sum().backward() + torch.cuda.synchronize() + return [out, inp.grad] + [p.grad for p in block.parameters()] + + outputs_ref = run(sequential_linear) + outputs = run(grouped_linear) + + ran = {phase for phase, ok, _ in dispatches if ok} + fell_back = sorted({phase for phase, ok, _ in dispatches if not ok}) + assert not fell_back and ran == {"fwd", "dgrad", "wgrad"}, ( + f"FlyDSL grouped GEMM ran {sorted(ran)}, fell back in {fell_back}. " + "See the [FLYDSL WARNING] for the fallback reason." + ) + for phase, _, scales in dispatches: + assert scales != {1.0} and len(scales) > 1, ( + f"{phase}: the compared step used scales {sorted(scales)}; it must exercise " + "non-unit, per-expert FP8 scales." + ) + for o, o_ref in zip(outputs, outputs_ref): + torch.testing.assert_close(o, o_ref, rtol=8e-3, atol=1e-3) + + @pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm-only grouped GEMM backend") @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_grouped_linear_cuda_graph_hipkittens_mxfp8(monkeypatch, capfd): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index aa20712c7..6e6390ec6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" +"""FlyDSL GEMM kernels: dense BF16/FP16/FP32/FP8/MXFP8, grouped BF16/FP16/FP8/MXFP8.""" # Supported flydsl release series. Single source of truth for the version # requirement, since flydsl is user-installed (not a declared dependency). A @@ -41,9 +41,10 @@ def _check_flydsl_version() -> None: _check_flydsl_version() from .exceptions import FlyDSLUnsupportedError -from .gemm_wrappers import te_generic_gemm_flydsl +from .gemm_wrappers import general_grouped_gemm_flydsl, te_generic_gemm_flydsl __all__ = [ "FlyDSLUnsupportedError", + "general_grouped_gemm_flydsl", "te_generic_gemm_flydsl", ] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 1b7f9250e..9afda422c 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -27,7 +27,12 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + placeholder_arg, + require_block_tiling, + require_launch_size, +) from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -198,10 +203,11 @@ def kernel_gemm( by_n = pid_n * BLOCK_N # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines # these values with i32 constants, so Index-typed coordinates would make @@ -231,7 +237,7 @@ def kernel_gemm( # Per-CTA tile base in C elements, folded into each store's linear # coordinate below (matching the scale-load addressing on this build; # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + c_tile_base_elems = bx_m_idx * fx.Int64(c_n) + by_n_idx gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) @@ -366,15 +372,15 @@ def hot_loop_scheduler_q_prefetch_4n(): def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one # 128x64 half-page (16 KiB). Each half has its own LDS base. - global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index( + global_base = (bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index( + global_base = (by_n_idx + fx.Int64(subtile * (BLOCK_N // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -403,8 +409,8 @@ def load_frag_at_byte_base(lds_page, row_byte_base): def load_b_frag(lds_b, local_row, half): # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K * ELEM_BYTES)) + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Int64(BLOCK_K * ELEM_BYTES)) def _acc_idx(subtile_id, mi, ni): return ( @@ -491,10 +497,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): mi = local_idx // MFMA_N_PER_SUBTILE ni = local_idx % MFMA_N_PER_SUBTILE - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 # Bias depends only on the output-feature (N) coordinate, so read it # once per column (global index by_n_idx + col) and reuse across the @@ -503,8 +509,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): bias_value = load_bias(by_n_idx + col) for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col # Epilogue stages run on the fp32 accumulator; output is fp32 # so there is no dtype narrowing: @@ -546,8 +552,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): zero_pinned_accumulators() def load_b_subtile_ni_regs(lds_b, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 return load_b_frag(lds_b, b_row_addr, sn) def load_b_subtile_regs(lds_b, sn): @@ -559,10 +565,10 @@ def load_b_subtile_regs(lds_b, sn): ) def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K * ELEM_BYTES) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Int64(BLOCK_K * ELEM_BYTES) return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) def load_a_subtile_mi_regs(lds_a, sm, mi): @@ -605,7 +611,7 @@ def hk_one_k_with_refill( b13 = load_b_subtile_ni_regs(cur_b, 1, 3) # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + k_refill = fx.Int64((k128 + 2) * BLOCK_K) # Q0: interleave the current tile's A-bottom LDS reads with K+2 # refills and Q0 compute. Compute is K16-half-major across all 16 @@ -864,15 +870,15 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): store_acc_vector_for_logical_idx(drain_acc_idx, acc) # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) @@ -989,27 +995,16 @@ def fp32_matmul( The optimized core streams both operands with K contiguous and consumes B as [N, K] directly, so this adapter performs no internal transpose. The wrapper is responsible for normalizing NN/NT layouts into this contract. + + Only the checks ``doGemm`` does not repeat are performed here: operand rank, + device agreement, and output contiguity. Inner-dim match, operand/output + dtype, and C shape are validated by ``doGemm`` against the runtime geometry, + so they are not duplicated in this shim. """ if a.ndim != 2 or b.ndim != 2: raise ValueError( f"FlyDSL FP32 TN expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" ) - - m, k = a.shape - n, kb = b.shape - if kb != k: - raise ValueError(f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}") - if a.dtype != torch.float32 or b.dtype != torch.float32: - raise TypeError( - "FlyDSL FP32 GEMM expects both operands to have torch.float32 dtype, " - f"got {a.dtype} and {b.dtype}" - ) - if tuple(c.shape) != (m, n): - raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype != torch.float32: - raise FlyDSLUnsupportedError( - f"The current FlyDSL FP32 kernel stores torch.float32 output, got {c.dtype}" - ) if a.device != b.device or a.device != c.device: raise ValueError( f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" @@ -1090,13 +1085,24 @@ def doGemm( B_arg = B.contiguous().view(torch.uint8).view(-1) C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # The placeholders are never dereferenced by the kernel, so reuse one cached + # buffer per (device, dtype) instead of re-allocating each call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.contiguous().view(-1) else: - Aux_arg = torch.zeros(1, dtype=torch.float32, device=A.device) - launch = _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue) - launch(A_arg, B_arg, C_arg, Bias_arg, Aux_arg, M_runtime, N_runtime, stream=stream) + Aux_arg = placeholder_arg(A.device, torch.float32) + + launch_key = ("fp32_gemm", int(K_runtime), bool(use_xcd_remap), epilogue) + launch_args = (A_arg, B_arg, C_arg, Bias_arg, Aux_arg, M_runtime, N_runtime, stream) + # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are + # runtime scalars, so one compiled launcher per launch_key serves all shapes. + compiled_launch( + launch_key, + _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue), + launch_args, + )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 7a9427f95..95468104d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -25,8 +25,15 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + get_num_cus, + placeholder_arg, + require_block_tiling, + require_launch_size, +) # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( @@ -95,6 +102,12 @@ def _compile_kernel( output_dtype: torch.dtype, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, + weight_ptrs: bool = False, + act_ptrs: bool = False, + scale_ptrs: bool = False, ): """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. @@ -194,6 +207,9 @@ class SharedStorage: b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + # Persistent grouped path: one-int32 scratch to broadcast the atomically + # claimed queue slot from thread 0 to the whole workgroup. + claim_slot: fx.Array[fx.Int32, 1, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -206,6 +222,12 @@ def kernel_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -219,741 +241,914 @@ def kernel_gemm( gB = make_fp8_buffer_tensor(B, b_f8_ir_t) a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) - a_scale_rsrc = fx.rocdl.make_buffer_tensor(A_scale_inv, max_size=True) - b_scale_rsrc = fx.rocdl.make_buffer_tensor(B_scale_inv, max_size=True) - a_scale_div = fx.logical_divide(a_scale_rsrc, fx.make_layout(1, 1)) - b_scale_div = fx.logical_divide(b_scale_rsrc, fx.make_layout(1, 1)) + # scale_ptrs (grouped only): A_scale_inv / B_scale_inv are int64[G] arrays of + # per-expert scale_inv pointers, read in place per tile instead of from a + # host-gathered length-G vector. + if const_expr(not scale_ptrs): + a_scale_rsrc = fx.rocdl.make_buffer_tensor(A_scale_inv, max_size=True) + b_scale_rsrc = fx.rocdl.make_buffer_tensor(B_scale_inv, max_size=True) + a_scale_div = fx.logical_divide(a_scale_rsrc, fx.make_layout(1, 1)) + b_scale_div = fx.logical_divide(b_scale_rsrc, fx.make_layout(1, 1)) scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - def _load_scale(scale_div): + def _load_scale(scale_div, index): + # Dense passes a 1-element scale buffer (index 0). Grouped passes a + # length-num_groups vector indexed by the tile's expert, so each tile + # recomputes its own per-expert output_scale inside process_tile. reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(scale_ld_atom, fx.slice(scale_div, (None, fx.Int32(0))), reg) + fx.copy(scale_ld_atom, fx.slice(scale_div, (None, index)), reg) return fx.memref_load_vec(reg)[0] - output_scale = _load_scale(a_scale_div) * _load_scale(b_scale_div) tx = gpu.thread_id("x") + tx_i32 = fx.Int32(tx) num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - if const_expr(use_xcd_remap): - pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) - else: - pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) - - bx_m = pid_m * BLOCK_M - by_n = pid_n * BLOCK_N - - # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines - # these values with i32 constants, so Index-typed coordinates would make - # arith.addi receive mixed operand types. - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # The utility mapping is identical to the previous manual staging: - # each step contributes one contiguous 16-byte vector per thread, while - # the global K coordinate is XOR-unswizzled for the physical LDS slot. - gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) - gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) - a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) - b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) - s2r = S2RLoader(fx.Int32(0), 1) - - layout_lane16 = fx.make_layout((4, 16), (16, 1)) - coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) - lane_div_16 = fx.get(coord_lane16, 0) - lane_mod_16 = fx.get(coord_lane16, 1) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - # Bias is a length-N fp32 vector indexed by the global output-feature - # (N) coordinate and broadcast across the M/token rows. const_expr folds - # this compile-time flag at trace time so the setup (and load_bias) are - # inlined into the kernel scope with no runtime dispatch branch. - if const_expr(has_bias): - gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) - bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) - bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - - def load_bias(col): - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) - return fx.memref_load_vec(reg)[0] - - # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N - # output so the backward pass can recompute the GELU gradient. Same tile - # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. - if const_expr(has_gelu): - gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) - aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) - aux_store_atom = fx.make_copy_atom( + # ``grouped`` gates the shared descriptor queue / work-steal driver; + # const_expr folds it out of the dense compiled kernel. + grouped = group or wgrad + # Grouped drain bound: the exact number of descriptors the host wrote to + # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. + num_tiles = num_work + + # weight_ptrs: weights are G independently-allocated per-expert slabs, + # addressed from a per-expert base pointer read from ``b_data_ptrs`` (reusing + # gB's static layout) instead of an expert-shifted B row. FP8 is TN-only, so + # both grouped forward and dgrad (dgrad feeds host-transposed weights through + # this same TN kernel) set this compile-time flag. wgrad sets it only with + # act_ptrs, for its per-expert X^T operand. + if const_expr(weight_ptrs): + b_ptr_layout = fx.get_layout(gB) + # act_ptrs: the per-expert activation operand (A; and in wgrad also B, via + # b_data_ptrs) stays in its own per-expert allocation instead of a host-side + # concat. The host passes each expert's base pointer pre-shifted by that + # expert's packed offset, so the packed addressing below is unchanged; only the + # buffer base moves per tile. + if const_expr(act_ptrs): + a_ptr_layout = fx.get_layout(gA) + + if const_expr(grouped): + wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) + wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + claim_iter, claim_view = grouped_utils.make_claim_broadcast(lds.claim_slot) + + def _claim_next_slot(): + if tx_i32 == fx.Int32(0): + old = llvm.AtomicRMWOp( + llvm.AtomicBinOp.add, + grouped_utils.global_i32_ptr(work_counter, fx.Int32(0)), + fx.as_ir_value(fx.Int32(1)), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + fx.ptr_store(fx.Vector.from_elements([fx.Int32(old)], fx.Int32), claim_iter) + fx.barrier() + return fx.Vector(claim_view.load())[0] + + def _desc_field(slot, field): + return grouped_utils.read_desc_field(wq_div, desc_ld_atom, slot, field) + + def process_tile(pid_m, pid_n, expert): + # Per-expert scalar scale: dense reads index 0 (length-1 buffer); grouped + # reads scale[expert] from the length-num_groups vector, or through the + # expert's scale_inv pointer (scale_ptrs). + if const_expr(scale_ptrs): + a_scale = fx.Float32(grouped_utils.load_f32_at_ptr(A_scale_inv, expert)) + b_scale = fx.Float32(grouped_utils.load_f32_at_ptr(B_scale_inv, expert)) + else: + if const_expr(grouped): + scale_idx = fx.Int32(expert) + else: + scale_idx = fx.Int32(0) + a_scale = _load_scale(a_scale_div, scale_idx) + b_scale = _load_scale(b_scale_div, scale_idx) + output_scale = a_scale * b_scale + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) + + # Per-expert offsets fold the expert index into the packed operands / + # output (const_expr-zeroed for dense). FP8 is TN-NORMAL-read, so the + # contraction axis (K) is the trailing dim and is NEVER packed; experts + # always stack on an OUTPUT/feature axis: + # + # fwd/dgrad: A is per-token (absolute M rows, no offset); B (weights, + # or W^T for dgrad) is packed on its N-row by expert + # (by_n_b_idx = expert*c_n + n). Output per-token (no store offset). + # + # wgrad (dW_e = dY_e^T @ X_e): host pre-transposes both operands to + # [feature, m_pad] and token-pads the trailing (contraction) axis, so + # A = dY^T packed [G*N_out, m_pad] on the M-row (a_expert_m_off = + # expert*c_m), B = X^T packed [G*K_in, m_pad] on the N-row + # (by_n_b_idx = expert*c_n + n). Output dW packed [G*N_out, K_in] on + # the M-row (store_m_expert_off = expert*c_m). + if const_expr(wgrad): + by_n_b_idx = fx.Int64(expert * fx.Int32(c_n) + by_n) + a_expert_m_off = fx.Int64(expert * fx.Int32(c_m)) + store_m_expert_off = fx.Int64(expert * fx.Int32(c_m)) + elif const_expr(group): + by_n_b_idx = fx.Int64(expert * fx.Int32(c_n) + by_n) + a_expert_m_off = fx.Int64(0) + store_m_expert_off = fx.Int64(0) + else: + by_n_b_idx = by_n_idx + a_expert_m_off = fx.Int64(0) + store_m_expert_off = fx.Int64(0) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + # Issue both per-expert base-pointer loads before either is made uniform, so a + # tile waits on one memory round trip rather than two. + if const_expr(act_ptrs): + a_base_raw = grouped_utils.read_ptr_at(a_data_ptrs, expert) + if const_expr(weight_ptrs): + b_base_raw = grouped_utils.read_ptr_at(b_data_ptrs, expert) + if const_expr(act_ptrs): + a_base_i64 = grouped_utils.make_uniform_i64(a_base_raw) + a_div_tile = grouped_utils.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) + else: + a_div_tile = a_div + a_g2s = G2SLoader(a_div_tile, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + # Grouped-TN forward B-data: rebase onto this tile's expert weight slab and + # read with a LOCAL column (base selects the expert). Other paths keep the + # shared packed b_div / by_n_b_idx. + if const_expr(weight_ptrs): + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + # wgrad bases are pre-shifted by the packed expert offset (see act_ptrs), + # so keep the packed column there; forward/dgrad weight bases are the + # real per-expert slabs, read with a local column. + if const_expr(wgrad): + b_data_col = by_n_b_idx + else: + b_data_col = by_n_idx + else: + b_div_tile = b_div + b_data_col = by_n_b_idx + b_g2s = G2SLoader(b_div_tile, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + # store_m_expert_off (grouped wgrad only) repacks the output row to + # this expert's block; 0 for dense and grouped fwd/dgrad. + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Int64(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), output_fx_dtype, ) - def gelu_tanh(x): - # tanh-approx GELU (matches PyTorch approximate='tanh' and the - # FlyDSL preshuffle reference), expressed through a non-positive - # exponent so exp() cannot overflow: - # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = x * x * x - y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * x * (numerator * (one_f32 / denom)) - - PIN_ACC_BASE = 0 - - def _reg_list(prefix, start, end): - return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) - - def reserve_pinned_accumulators(): - # Reserve a fixed physical AGPR bank for all accumulators. In the - # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, - # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator - # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the - # scaled MFMA accumulation in place and avoids those transfers and spills. - # - # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, - # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. - clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) - llvm.InlineAsmOp( - None, - [], - "", - clobbers, - has_side_effects=True, - ) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy32b() + if output_element_bytes == 4 + else fx.rocdl.BufferCopy16b() + ), + output_fx_dtype, + ) - def zero_pinned_accumulators(): - for ai in range_constexpr(ACCS_PER_WAVE * 4): + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) llvm.InlineAsmOp( None, [], - f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", - f"~{{a{PIN_ACC_BASE + ai}}}", + "", + clobbers, has_side_effects=True, ) - def _inline_asm_i32(asm_string, constraints, operands=None): - op = llvm.InlineAsmOp( - T.i32, - operands or [], - asm_string, - constraints, - has_side_effects=True, - ) - return _one_i32_result(op) - - def _one_i32_result(op): - # Accept the result attribute names exposed by the supported MLIR Python bindings. - return getattr(op, "result", getattr(op, "res", op.results[0])) - - def read_pinned_accumulator(acc_idx): - acc_pin = PIN_ACC_BASE + acc_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def read_physical_accumulator_slot(slot_idx): - acc_pin = PIN_ACC_BASE + slot_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def hot_loop_scheduler_q_refill_2n(): - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_mfma(2) - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q0_refill_a1_2n(): - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_dsrd(1) - rocdl.sched_mfma(2) - rocdl.sched_barrier(0) + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x128 half-page (16 KiB). Each half has its own LDS base. + # a_expert_m_off shifts the A row to this tile's expert block (grouped + # wgrad, where A = dY^T is packed [G*N_out, m_pad] on the M-row); 0 + # for dense and grouped fwd/dgrad (A is per-token, absolute rows). + global_base = ( + bx_m_idx + a_expert_m_off + fx.Int64(subtile * (BLOCK_M // 2)) + ) * fx.Int64(K) + k_base + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # Grouped-TN forward selects the expert via the rebased per-expert B + # view (b_g2s) with a LOCAL column (b_data_col). Other paths: by_n_b_idx + # shifts the B row to the expert block (grouped dgrad weights, grouped + # wgrad X^T); reduces to dense by_n for the dense path. + b_row = b_data_col + fx.Int64(subtile * (BLOCK_N // 2)) + global_base = b_row * fx.Int64(K) + k_base + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Int64(BLOCK_K)) + + def _acc_idx(subtile_id, mi, ni): + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + mi * MFMA_N_PER_SUBTILE + + ni + ) - def hot_loop_scheduler_q_prefetch_4n(): - for _ in range_constexpr(8): - rocdl.sched_dsrd(2) - rocdl.sched_mfma(4) - rocdl.sched_barrier(0) + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + "v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", + has_side_effects=True, + ) - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one - # 128x128 half-page (16 KiB). Each half has its own LDS base. - global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base - a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base - b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_a_subtile(k_base, subtile, lds_a): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - - def stage_b_subtile(k_base, subtile, lds_b): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) - - def load_frag_half_at_byte_base(lds_page, row_byte_base, half): - # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. - # Keeping the halves separate allows steady-state Q0 to schedule one - # A-bottom ds_read_b128 in each refill/MFMA chunk. - k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 - return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) - - def pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_frag_at_byte_base(lds_page, row_byte_base): - # Default complete-fragment path used outside the dedicated Q0 schedule. - x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) - x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) - return pack_frag_halves(x0, x1) - - def load_b_frag(lds_b, local_row, half): - # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) - - def _acc_idx(subtile_id, mi, ni): - return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni - ) + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + "v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + "$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", + has_side_effects=True, + ) - def pinned_mfma(acc_idx, a_frag, b_frag): - """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" - acc_pin = PIN_ACC_BASE + acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - arith._to_raw(a_frag), - arith._to_raw(b_frag), - ], - "v_mfma_f32_16x16x128_f8f6f4 " - f"a[{acc_pin}:{acc_pin + 3}], " - "$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", - has_side_effects=True, - ) + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) - def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): - """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" - dst_pin = PIN_ACC_BASE + dst_slot * 4 - old_pin = PIN_ACC_BASE + old_acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - arith._to_raw(a_frag), - arith._to_raw(b_frag), - ], - "v_mfma_f32_16x16x128_f8f6f4 " - f"a[{dst_pin}:{dst_pin + 3}], " - "$0, $1, " - f"a[{old_pin}:{old_pin + 3}] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", - has_side_effects=True, - ) + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) - def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - pinned_mfma(acc_base + 2, a_frag, b2) - pinned_mfma(acc_base + 3, a_frag, b3) - - def mfma_2n(acc_base, a_frag, b0, b1): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - - def store_acc_vector_for_logical_idx(logical_acc_idx, acc): - subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - - # Bias depends only on the output-feature (N) coordinate, so read it - # once per column (global index by_n_idx + col) and reuse across the - # four M rows below. - if const_expr(has_bias): - bias_value = load_bias(by_n_idx + col) + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 - # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing. Bias is added after the tensor-wise - # output scale, matching the reference epilogue ordering; GELU_AUX - # then saves the pre-activation and applies GELU to the C value. - value = Vec(acc)[ii] * output_scale + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. if const_expr(has_bias): - value = value + bias_value + bias_value = load_bias(by_n_idx + col) + + for ii in range_constexpr(4): + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing. Bias is added after the tensor-wise + # output scale, matching the reference epilogue ordering; GELU_AUX + # then saves the pre-activation and applies GELU to the C value. + value = Vec(acc)[ii] * output_scale + if const_expr(has_bias): + value = value + bias_value + + if const_expr(has_gelu): + aux_val = value + if output_dtype != torch.float32: + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) - if const_expr(has_gelu): - aux_val = value if output_dtype != torch.float32: - aux_val = aux_val.to(output_fx_dtype) - aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) - fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) - value = gelu_tanh(value) - - if output_dtype != torch.float32: - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) - - # Explicit register coordinates for HK-style four-quadrant mapping. - # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions - # inside each 128x128 quadrant: - # cA: (warp_m, warp_n) - # cB: (warp_m, warp_n + 2) - # cC: (warp_m + 2, warp_n) - # cD: (warp_m + 2, warp_n + 2) - reg_k_col0 = lane_div_16 * 16 - reg_k_col1 = 64 + lane_div_16 * 16 - - # Every fragment row differs only by multiples of 16, so row % 16 is - # always lane_mod_16. Hoist the logical->physical XOR mapping once. - _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) - _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) - - reg_subtile_m_idx0 = wave_id // 2 - reg_subtile_n_idx0 = wave_id % 2 - - reserve_pinned_accumulators() - zero_pinned_accumulators() - - def load_b_subtile_ni_regs(lds_b, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 - return load_b_frag(lds_b, b_row_addr, sn) - - def load_b_subtile_regs(lds_b, sn): - return ( - load_b_subtile_ni_regs(lds_b, sn, 0), - load_b_subtile_ni_regs(lds_b, sn, 1), - load_b_subtile_ni_regs(lds_b, sn, 2), - load_b_subtile_ni_regs(lds_b, sn, 3), - ) - - def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K) - return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) - - def load_a_subtile_mi_regs(lds_a, sm, mi): - x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) - x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) - return pack_frag_halves(x0, x1) - - def load_a_subtile_regs(lds_a, sm): - return ( - load_a_subtile_mi_regs(lds_a, sm, 0), - load_a_subtile_mi_regs(lds_a, sm, 1), - load_a_subtile_mi_regs(lds_a, sm, 2), - load_a_subtile_mi_regs(lds_a, sm, 3), - ) - - def hk_one_k_with_refill( - k128, - cur_a, - cur_b, - next_a, - next_b, - refill_a, - refill_b, - a0_regs, - b0_regs, - ): - - # Wait only far enough for the current page; the next-page refill may remain in flight. - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) - - # A-top and B-left are both carried as complete 64-row register tiles, - # so their LDS half-pages can be refilled immediately. - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs - - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - - # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) - - # Q0: interleave the current tile's A-bottom LDS reads with K+2 - # refills and Q0 compute. Each complete A-bottom fragment is assembled - # from two independently scheduled K64 halves. - rocdl.sched_barrier(0) - a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a) - mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) + value = value.to(output_fx_dtype) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = ( + subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 + ) + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) - a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) - stage_b_subtile_pass(k_refill, 0, 0, refill_b) - mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = ( + subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + ) + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Int64(BLOCK_K) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) - a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) - stage_a_subtile_pass(k_refill, 0, 1, refill_a) - mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): - a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) - stage_b_subtile_pass(k_refill, 0, 1, refill_b) - mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) + # Wait only far enough for the current page; the next-page refill may remain in flight. + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) - stage_a_subtile_pass(k_refill, 0, 2, refill_a) - mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) - stage_b_subtile_pass(k_refill, 0, 2, refill_b) - mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) - stage_a_subtile_pass(k_refill, 0, 3, refill_a) - mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Int64((k128 + 2) * BLOCK_K) - a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) - stage_b_subtile_pass(k_refill, 0, 3, refill_b) - mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) - hot_loop_scheduler_q0_refill_a1_2n() + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) - # Retire the eight distributed A-bottom LDS reads before K+2 refills - # overwrite the current page's A-bottom half-page. Keep this wait as - # late as possible to maximize read/compute overlap. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) - a10 = pack_frag_halves(a10_x0, a10_x1) - a11 = pack_frag_halves(a11_x0, a11_x1) - a12 = pack_frag_halves(a12_x0, a12_x1) - a13 = pack_frag_halves(a13_x0, a13_x1) + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) - rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 0, refill_b) - mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) - stage_a_subtile_pass(k_refill, 1, 0, refill_a) - mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) - stage_b_subtile_pass(k_refill, 1, 1, refill_b) - mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) - stage_a_subtile_pass(k_refill, 1, 1, refill_a) - mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) - stage_b_subtile_pass(k_refill, 1, 2, refill_b) - mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + hot_loop_scheduler_q0_refill_a1_2n() - stage_a_subtile_pass(k_refill, 1, 2, refill_a) - mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 3, refill_b) - mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) - stage_a_subtile_pass(k_refill, 1, 3, refill_a) - mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) - hot_loop_scheduler_q_refill_2n() + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) - # Leave exactly the K+2 refill and scale loads outstanding. The following - # LDS reads consume the already-ready next page, not the page being refilled. - rocdl.sched_barrier(0) - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + hot_loop_scheduler_q_refill_2n() - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) - hot_loop_scheduler_q_prefetch_4n() + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) - return next_a0_regs, next_b0_regs + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) - def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) - mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + hot_loop_scheduler_q_prefetch_4n() - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + return next_a0_regs, next_b0_regs - rocdl.sched_barrier(0) - barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + rocdl.sched_barrier(0) + barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) - hot_loop_scheduler_q_prefetch_4n() + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) - return next_a0_regs, next_b0_regs + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) - def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): - barrier(vmcnt=0, lgkmcnt=0) + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) - # Materialize the remaining final-page A/B fragments once. The - # subsequent schedule is entirely register/AGPR traffic. - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + return next_a0_regs, next_b0_regs - a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) - b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + barrier(vmcnt=0, lgkmcnt=0) - # Rolling final-page epilogue. - # - # Finalize accumulators in their own physical AGPR slots, but delay - # each AGPR read/store until several independent final MFMAs have - # been issued. - # - # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, - # MFMA 4, drain 1, MFMA 5, drain 2, ... - # - # The buffer stores are only issued here; they may remain in flight - # while later MFMAs and accumulator drains continue. - FINAL_EPILOGUE_DEPTH = 4 - pending = [] - - for old_acc_idx in range_constexpr(ACCS_PER_WAVE): - subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi - b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni - - # Final MFMA remains in-place. The logical accumulator's own - # AGPR slot is unique and cannot conflict with another pending - # result, so no ad-hoc physical-slot permutation is needed. - pinned_final_mfma( - old_acc_idx, - old_acc_idx, - a_frags[a_frag_idx], - b_frags[b_frag_idx], - ) - pending.append(old_acc_idx) + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - # Drain the oldest completed result only after enough newer - # independent MFMAs have supplied the MFMA->AGPR-read spacing. - if len(pending) == FINAL_EPILOGUE_DEPTH: - drain_acc_idx = pending.pop(0) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: acc = read_physical_accumulator_slot(drain_acc_idx) store_acc_vector_for_logical_idx(drain_acc_idx, acc) - # Flush the final results after all final-page MFMAs have issued. - for drain_acc_idx in pending: - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) - # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) - - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) - - a0_regs = load_a_subtile_regs(lds_a0, 0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) + a0_regs = load_a_subtile_regs(lds_a0, 0) - b0_regs = load_b_subtile_regs(lds_b0, 0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - # Main HK loop: exactly one logical K128 per iteration. - # Even k consumes and refills LDS0; odd k does the same for LDS1. - for k128 in range_constexpr(NUM_K_TILES - 2): - if (k128 % 2) == 0: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a0, lds_b0, lds_a1, lds_b1, - lds_a0, - lds_b0, a0_regs, b0_regs, ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) else: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a1, lds_b1, lds_a0, lds_b0, - lds_a1, - lds_b1, a0_regs, b0_regs, ) - - # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch - # to prepare A-top/B-left for the final tile, but performs no K+2 refill. - if (NUM_K_TILES % 2) == 0: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a0, - lds_b0, - lds_a1, - lds_b1, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + # Dispatch. Dense runs one XCD-swizzled tile. Grouped runs a persistent + # loop over precomputed (pid_m, pid_n, expert) descriptors, selecting the + # schedule at compile time (``sched``): "worksteal" atomically claims the + # next global slot each iteration; "static" owns a fixed grid-stride slice. + def _run_grouped_tile(slot): + pid_m = _desc_field(slot, fx.Int32(0)) + pid_n = _desc_field(slot, fx.Int32(1)) + expert = _desc_field(slot, fx.Int32(2)) + process_tile(pid_m, pid_n, expert) + + if const_expr(grouped): + if const_expr(sched == "static"): + stride = fx.Int32(gpu.grid_dim.x) + slot = fx.Int32(gpu.block_id("x")) + while slot < num_tiles: + _run_grouped_tile(slot) + slot = slot + stride + else: + active = fx.Int32(1) == fx.Int32(1) + while active: + slot = _claim_next_slot() + in_range = slot < num_tiles + if in_range: + _run_grouped_tile(slot) + active = in_range else: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a1, - lds_b1, - lds_a0, - lds_b0, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + if const_expr(use_xcd_remap): + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + process_tile(pid_m, pid_n, fx.Int32(0)) @flyc.jit def launch_gemm( @@ -966,10 +1161,20 @@ def launch_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, stream: fx.Stream = fx.Stream(None), ): - # The integration only dispatches aligned shapes; no partial-tile masking exists. + # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU + # persistent grid that drains the descriptor queue. Placeholder work + # tensors for the dense path are allocated host-side in ``doGemm``. grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(group or wgrad): + grid_x = get_num_cus() kernel_gemm( A, B, @@ -980,6 +1185,12 @@ def launch_gemm( Aux, c_m, c_n, + num_groups, + work_queue, + work_counter, + b_data_ptrs, + a_data_ptrs, + num_work, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -994,6 +1205,12 @@ def _cached_launch( output_dtype: torch.dtype, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, + weight_ptrs: bool = False, + act_ptrs: bool = False, + scale_ptrs: bool = False, ): return _compile_kernel( K, @@ -1002,6 +1219,12 @@ def _cached_launch( output_dtype, use_xcd_remap=use_xcd_remap, epilogue=epilogue, + group=group, + sched=sched, + wgrad=wgrad, + weight_ptrs=weight_ptrs, + act_ptrs=act_ptrs, + scale_ptrs=scale_ptrs, ) @@ -1029,54 +1252,13 @@ def fp8_matmul( The wrapper is responsible for adapting TE's logical [K, N] operand into the core's row-major [N, K] representation. No operand swap or transpose is performed in this module. - """ - if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): - raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") - - if a.ndim != 2 or b.ndim != 2: - raise ValueError( - f"FlyDSL FP8 TN expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" - ) - - supported_fp8_dtypes = ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ) - if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: - raise TypeError( - f"FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, got A={a.dtype} and B={b.dtype}" - ) - - m, k = a.shape - n, kb = b.shape - if kb != k: - raise ValueError(f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}") - - for name, scale in ( - ("A_scale_inv", a_scale_inv), - ("B_scale_inv", b_scale_inv), - ): - if not isinstance(scale, torch.Tensor): - raise TypeError(f"{name} must be a torch.Tensor") - if scale.dtype != torch.float32 or scale.numel() != 1: - raise TypeError( - f"{name} must contain exactly one FP32 value, got " - f"dtype={scale.dtype}, shape={tuple(scale.shape)}" - ) - - if tuple(c.shape) != (m, n): - raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 supports only float16, bfloat16, and float32 outputs, got {c.dtype}" - ) - if not c.is_contiguous(): - raise FlyDSLUnsupportedError("FlyDSL FP8 requires contiguous output storage") - - tensors = (a, b, a_scale_inv, b_scale_inv, c) - if any(t.device != a.device for t in tensors[1:]): - raise ValueError("A, B, inverse scales, and C must be on the same device") + ``_run_fp8`` (the sole caller) has already validated operand rank, FP8 + dtypes, inverse-scale dtype/numel, and contiguity; ``doGemm`` re-validates + dtypes, inner-dim match, scale numel, and C shape/dtype against the runtime + geometry. This shim therefore just forwards, rather than repeating those + checks a third time (``doGemm`` derives M/N/K from the operands itself). + """ doGemm( a, b, @@ -1101,10 +1283,35 @@ def doGemm( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: torch.Tensor = None, + work_counter: torch.Tensor = None, + b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, + num_work: int = 0, + sched: str = "worksteal", + wgrad: bool = False, + m: int = None, + n: int = None, + k: int = None, + a_shape=None, + b_shape=None, ): - """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" - M_runtime, K_runtime = A.shape - N_runtime, Kb_runtime = B.shape + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales. + + Dense derives M/N/K from operand shapes. Grouped/wgrad pass explicit per-expert + ``m``/``n``/``k`` (the packed operands do not reveal the per-expert geometry), + a length-``num_groups`` scale vector for each operand, and the descriptor queue. + With ``a_data_ptrs`` (``b_data_ptrs`` in wgrad) the operand is one representative + expert and ``a_shape`` (``b_shape``) the packed operand shape. + """ + if group or wgrad: + M_runtime, N_runtime, K_runtime = int(m), int(n), int(k) + else: + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" supported_fp8_dtypes = ( torch.float8_e4m3fn, torch.float8_e5m2, @@ -1116,7 +1323,6 @@ def doGemm( torch.bfloat16, torch.float32, ), f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" require_block_tiling( M_runtime, N_runtime, @@ -1126,13 +1332,29 @@ def doGemm( block_k=_BLOCK_K, label="FP8 GEMM", ) - require_launch_size("FP8 GEMM", ("A", A), ("B", B), ("C", C)) - assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 - assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 - assert C.shape == ( - M_runtime, - N_runtime, - ), f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + # With per-expert pointers, A (and wgrad's B) is one representative expert, but + # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit + # offsets, so the bound is on the packed extent. + a_rows, a_cols = a_shape if a_shape is not None else A.shape + b_rows, b_cols = b_shape if b_shape is not None else B.shape + require_launch_size( + "FP8 GEMM", + ("A", A, a_rows * a_cols * A.element_size()), + ("B", B, b_rows * b_cols * B.element_size()), + ("C", C), + ) + # Dense: one scalar inverse scale per operand. Grouped: a per-expert vector, or + # (scale_ptrs) an int64[G] array of per-expert scale_inv pointers read in place. + expected_scale_numel = num_groups if (group or wgrad) else 1 + scale_ptrs = A_scale_inv.dtype == torch.int64 + expected_scale_dtype = torch.int64 if scale_ptrs else torch.float32 + assert not scale_ptrs or group or wgrad, "per-expert scale pointers are grouped-only" + assert A_scale_inv.dtype == expected_scale_dtype and A_scale_inv.numel() == expected_scale_numel + assert B_scale_inv.dtype == expected_scale_dtype and B_scale_inv.numel() == expected_scale_numel + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is + # the unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + assert tuple(C.shape) == expected_c, f"C shape {tuple(C.shape)} != {expected_c}" if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported FP8 epilogue: {epilogue}") @@ -1166,7 +1388,7 @@ def doGemm( raise ValueError(f"FP8 epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = torch.cuda.current_stream() + stream = grouped_utils.current_stream_handle(A.device) A_arg = A.view(torch.uint8).contiguous().view(-1) B_arg = B.view(torch.uint8).contiguous().view(-1) @@ -1174,24 +1396,65 @@ def doGemm( A_scale_arg = A_scale_inv.contiguous().view(-1) B_scale_arg = B_scale_inv.contiguous().view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # The placeholders are never dereferenced by the kernel, so reuse one cached + # buffer per (device, dtype) instead of re-allocating each call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + Aux_arg = placeholder_arg(A.device, C.dtype) + + # Grouped runtime tensors are required by the kernel signature; the dense + # compiled kernel never dereferences them but still needs placeholder buffers. + if group or wgrad: + # num_work is the kernel's drain bound: it must be the queue's descriptor count. + if work_queue.numel() != 3 * int(num_work): + raise ValueError( + f"num_work={num_work} does not match work_queue of {work_queue.numel()} " + "int32 (3 per descriptor)" + ) + work_queue_arg = work_queue.contiguous().view(-1) + work_counter_arg = work_counter.contiguous().view(-1) + else: + work_queue_arg = placeholder_arg(A.device, torch.int32) + work_counter_arg = placeholder_arg(A.device, torch.int32) + + # Per-expert weight base pointers (grouped forward and dgrad -- both run on this + # TN kernel, dgrad over host-transposed weights): compile a kernel that reads each + # expert's weight slab from b_data_ptrs[g] instead of concatenating the G weights + # host-side. weight_ptrs is a compile-time flag keyed into the kernel cache; wgrad + # leaves it False and passes a placeholder never dereferenced. + weight_ptrs = b_data_ptrs is not None + if weight_ptrs: + b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) + else: + b_data_ptrs_arg = placeholder_arg(A.device, torch.int64) + # Per-expert activation base pointers (see act_ptrs in _compile_kernel). + act_ptrs = a_data_ptrs is not None + if act_ptrs: + a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) + else: + a_data_ptrs_arg = placeholder_arg(A.device, torch.int64) - launch = _cached_launch( + launch_key = ( + "fp8_gemm", int(K_runtime), A.dtype, B.dtype, C.dtype, bool(use_xcd_remap), epilogue, + group, + sched, + wgrad, + weight_ptrs, + act_ptrs, + scale_ptrs, ) - launch( + launch_args = ( A_arg, B_arg, C_arg, @@ -1201,5 +1464,32 @@ def doGemm( Aux_arg, M_runtime, N_runtime, - stream=stream, + num_groups, + work_queue_arg, + work_counter_arg, + b_data_ptrs_arg, + a_data_ptrs_arg, + int(num_work), + stream, ) + # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are + # runtime scalars, so one compiled launcher per launch_key serves all shapes. + compiled_launch( + launch_key, + _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + epilogue, + group, + sched, + wgrad, + weight_ptrs, + act_ptrs, + scale_ptrs, + ), + launch_args, + )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index f91aa6e04..a5996a37f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -12,12 +12,28 @@ G2S/S2R loaders, buffer-tensor makers) stays in the per-dtype modules. """ +import functools + +import torch + +import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import arith, const_expr, range_constexpr, rocdl from .exceptions import FlyDSLUnsupportedError +def get_num_cus(device=None) -> int: + """Number of CUs on ``device`` (default: the current device), cached per device index.""" + index = None if device is None else torch.device(device).index + return _num_cus_on(torch.cuda.current_device() if index is None else index) + + +@functools.lru_cache(maxsize=None) +def _num_cus_on(index: int) -> int: + return torch.cuda.get_device_properties(index).multi_processor_count + + def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_tiles=4): """Validate the host-side M/N/K tiling contract shared by every GEMM core. @@ -57,15 +73,17 @@ def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_til def require_launch_size(label, *tensors): """Reject operands whose byte size exceeds the int32 launch-argument limit. - ``tensors`` is an iterable of ``(name, tensor)`` pairs. Raises - ``FlyDSLUnsupportedError`` (so ``general_gemm`` falls back to the C++ - backend) for the first operand at or above ``_MAX_LAUNCH_BYTES``; skips - ``None`` tensors (e.g. an absent bias/aux). + ``tensors`` is an iterable of ``(name, tensor)`` or ``(name, tensor, nbytes)`` + entries. ``nbytes`` overrides the tensor's own size: a grouped operand read + through per-expert pointers is passed as one representative expert, while the + kernel addresses the whole packed operand. Raises ``FlyDSLUnsupportedError`` + (so ``general_gemm`` falls back to the C++ backend) for the first operand at or + above ``_MAX_LAUNCH_BYTES``; skips ``None`` tensors (e.g. an absent bias/aux). """ - for name, t in tensors: + for name, t, *override in tensors: if t is None: continue - nbytes = t.numel() * t.element_size() + nbytes = override[0] if override else t.numel() * t.element_size() if nbytes > _MAX_LAUNCH_BYTES: raise FlyDSLUnsupportedError( f"FlyDSL {label} operand {name} is {nbytes} bytes, which exceeds " @@ -73,6 +91,52 @@ def require_launch_size(label, *tensors): ) +@functools.lru_cache(maxsize=None) +def placeholder_arg(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """One cached 1-element buffer per (device, dtype) for an unused kernel argument. + + The FlyDSL kernels keep a uniform signature, so a dense launch still passes the + grouped/epilogue tensors (Bias/Aux/work_queue/work_counter/data-pointer arrays) + and a static-schedule grouped launch still passes the work counter. The kernel + never reads or writes those buffers, so one shared zero tensor per (device, dtype) + avoids a per-launch allocate+zero. + """ + return torch.zeros(1, dtype=dtype, device=device) + + +# Per-key cache of flyc.compile'd launchers, shared by every FlyDSL GEMM module. +# flydsl's JitFunction.__call__ rebuilds a full cache key (sig.bind + globals snapshot +# + arg-signature resolve) on every launch -- ~190us/call, which dominates wall time +# at small M where the kernel is short. flyc.compile returns a CompiledFunction whose +# __call__ skips all of that (ctypes pointer/scalar update + invoke, ~5us), baking +# constexpr values at compile time so one CompiledFunction per key serves every +# pointer/scalar/stream variation. M/N are runtime scalar args (not baked), so keys +# exclude them. Each caller's key starts with a module-unique tag (e.g. "fp8_gemm", +# "mx32_scale_pack"), so specializations from different modules never collide. +# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; +# these kernel modules never mutate their captured globals between launches, so it +# is safe here. +_COMPILED_LAUNCH_CACHE = {} + + +def compiled_launch(key, jit_fn, args): + """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. + + ``key`` uniquely identifies the compiled specialization (the same tuple the + JitFunction builder's lru_cache uses). ``args`` is the full positional arg list + for the first call; it is used once to compile, then the cached CompiledFunction + is reused positionally for all later calls with the same ABI. + """ + # A compiled launcher is built for the current device; key on it so a launch on + # another device compiles its own instead of reusing this one. + key = (key, torch.cuda.current_device()) + cf = _COMPILED_LAUNCH_CACHE.get(key) + if cf is None: + cf = flyc.compile(jit_fn, *args) + _COMPILED_LAUNCH_CACHE[key] = cf + return cf + + def cdiv(numer: int, denom: int) -> int: return (numer + denom - 1) // denom diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 4a1dbb09a..b99445d26 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -5,16 +5,23 @@ """TE entry points for the FlyDSL GEMM backend.""" import os +import warnings +from dataclasses import dataclass +from typing import Callable import torch import transformer_engine_torch as tex +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .half_prec_gemm import bf16_matmul, fp16_matmul +from .half_prec_gemm import _BLOCK_M as HALF_PREC_BLOCK_M, _BLOCK_N as HALF_PREC_BLOCK_N +from .half_prec_gemm import bf16_matmul, doGemm as half_prec_doGemm, fp16_matmul from .fp32_gemm import fp32_matmul -from .fp8_gemm import fp8_matmul -from .mxfp8_gemm import mxfp8_matmul +from .fp8_gemm import _BLOCK_M as FP8_BLOCK_M, _BLOCK_N as FP8_BLOCK_N +from .fp8_gemm import doGemm as fp8_doGemm, fp8_matmul +from .mxfp8_gemm import BLOCK_M as MXFP8_BLOCK_M, BLOCK_N as MXFP8_BLOCK_N +from .mxfp8_gemm import SCALE_GROUP_SIZE as MXFP8_SCALE_GROUP_SIZE, mxfp8_matmul # TODO: Some backend-independent GEMM wrapper utilities overlap with the # Triton GEMM backend (PR#667), including operand classification, logical @@ -1008,6 +1015,514 @@ def _run_mxfp8( return D, aux +# --------------------------------------------------------------------------------------- +# Grouped (MoE) GEMM: one persistent launch spans every expert. +# +# general_grouped_gemm_flydsl is a drop-in for cpp_extensions.general_grouped_gemm. The +# layout selects the phase (TN forward, NN dgrad, NT wgrad) and the operands' storage +# selects the kernel family. A _GroupedKernel describes the only parts that differ per +# family -- how an operand's payload and scale are read and how the kernel is launched -- +# while output validation, the descriptor queue, schedule choice, and the wgrad +# K-bucketing and scatter are shared by every family. +# --------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _GroupedKernel: + """How one FlyDSL kernel family reads grouped operands and launches.""" + + name: str + block_m: int + block_n: int + # tensor -> (payload, scale) whose contraction runs along the payload's last dim. + rows: Callable + # tensor -> (payload, scale) read the other way: dgrad's B and both wgrad operands. + cols: Callable + # (per-expert scales, device, k_pad) -> the kernel's scale argument. ``k_pad`` is + # the wgrad bucket's padded token count when the operands were padded, else None. + scale_arg: Callable + # (a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped) -> None + launch: Callable + # Kernel layout for dgrad: NN, or TN when cols() yields a materialized transpose. + dgrad_layout: str + # Axis of the wgrad operands (as returned by cols) that holds the tokens. + wgrad_token_axis: int + + +def _mxfp8_scale_arg(scales, device, k_pad): + # Per-expert E8M0 scales, packed by the kernel wrapper; a padded wgrad bucket pads + # them to the bucket's token count as well. + if k_pad is None: + return list(scales) + return grouped_utils.uniform_pad_concat(list(scales), k_pad // MXFP8_SCALE_GROUP_SIZE) + + +def _launch_mxfp8(a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped): + mxfp8_matmul(a, a_scale, b, b_scale, D, layout=layout, **grouped) + + +def _fp8_transpose(t): + """The materialized transpose of FP8 operand ``t`` (the TN-only kernel reads it).""" + if t._transpose is None or t._transpose_invalid: + raise FlyDSLUnsupportedError("FP8 operand has no valid transpose") + return t._transpose + + +def _launch_fp8(a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped): + # Tensor-wise FP8 is TN-only; dgrad and wgrad reach it through transposed operands. + fp8_doGemm(a, b, D, a_scale, b_scale, m=m, n=n, k=k, **grouped) + + +def _as_rows(t): + """``t`` as 2-D rows; GroupedLinear's per-expert tensors already are, so skip the view.""" + return t if t.dim() == 2 else t.reshape(-1, t.shape[-1]) + + +def _launch_half(a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped): + # BF16 and FP16 share this kernel; the MFMA opcode is selected per operand dtype. + half_prec_doGemm( + a, + b, + D, + layout=layout, + m=m, + n=n, + k=k, + input_dtype=a.dtype, + label="half-precision grouped", + **grouped, + ) + + +_MXFP8_GROUPED = _GroupedKernel( + name="MXFP8", + block_m=MXFP8_BLOCK_M, + block_n=MXFP8_BLOCK_N, + rows=lambda t: (t._rowwise_data, t._rowwise_scale_inv), + cols=lambda t: (t._columnwise_data, t._columnwise_scale_inv), + scale_arg=_mxfp8_scale_arg, + launch=_launch_mxfp8, + dgrad_layout="NN", + wgrad_token_axis=0, +) + +# One scalar inverse scale per expert, read in place through a pointer array. +_FP8_GROUPED = _GroupedKernel( + name="FP8", + block_m=FP8_BLOCK_M, + block_n=FP8_BLOCK_N, + rows=lambda t: (t._data, t._scale_inv), + cols=lambda t: (_fp8_transpose(t), t._scale_inv), + scale_arg=lambda scales, device, k_pad: grouped_utils.expert_ptrs(list(scales), device), + launch=_launch_fp8, + dgrad_layout="TN", + wgrad_token_axis=1, +) + +_HALF_GROUPED = _GroupedKernel( + name="half-precision", + block_m=HALF_PREC_BLOCK_M, + block_n=HALF_PREC_BLOCK_N, + rows=lambda t: (_as_rows(t), None), + cols=lambda t: (_as_rows(t), None), + scale_arg=lambda scales, device, k_pad: None, + launch=_launch_half, + dgrad_layout="NN", + wgrad_token_axis=0, +) + + +def _grouped_kernel_for(tensors): + """The kernel family for ``tensors``, which must all share one storage type.""" + from ...quantized_tensor import QuantizedTensorStorage + from ...tensor.storage import Float8TensorStorage, MXFP8TensorStorage + + if all(isinstance(t, MXFP8TensorStorage) for t in tensors): + return _MXFP8_GROUPED + if all(isinstance(t, Float8TensorStorage) for t in tensors): + return _FP8_GROUPED + if any(isinstance(t, QuantizedTensorStorage) for t in tensors): + raise FlyDSLUnsupportedError( + "FlyDSL grouped GEMM supports only MXFP8, tensor-wise FP8, and BF16/FP16 " + f"operands, got {sorted({type(t).__name__ for t in tensors})}" + ) + dtypes = {t.dtype for t in tensors} + if len(dtypes) == 1 and dtypes <= {torch.bfloat16, torch.float16}: + return _HALF_GROUPED + raise FlyDSLUnsupportedError( + f"FlyDSL half-precision grouped GEMM requires bf16/fp16 operands, got {dtypes}" + ) + + +def _as_kernel_dtype(data, t): + """Reinterpret TE's uint8 FP8 payload ``data`` of operand ``t`` as its FP8 dtype.""" + if data.dtype == torch.uint8: + return reinterpret_as_fp8_tensor(data, t._fp8_dtype) + return data + + +def _row_operand(datas, m_splits): + """Per-expert ``[m_e, K]`` row payloads as one grouped-kernel operand. + + Returns ``(data, data_ptrs)``: each expert's rows are read in place through base + pointers shifted to their packed row (``grouped_gemm_utils.stacked_operand``), so + the kernel sees the ``[M_total, K]`` layout without a concat. Falls back to + ``(concat_or_view_rows(datas), None)`` when the payloads are not row-major + ``[m_e, K]`` blocks. + """ + if all(d.dim() == 2 and d.size(0) == int(m) for d, m in zip(datas, m_splits)): + stacked = grouped_utils.stacked_operand(datas, datas[0].device) + if stacked is not None: + return stacked + return grouped_utils.concat_or_view_rows(datas), None + + +def _run_grouped(kernel, a_tensors, b_tensors, D, *, m_splits, layout): + """One packed forward or dgrad launch: ``D[rows of e] = A_e @ op(B_e)`` per expert. + + ``layout`` is the GEMM's TE layout. TN (forward): A = activations ``[M_total, K]``, + B = per-expert weights ``[N, K]``. NN (dgrad): A = grad_output ``[M_total, N_out]``, + B = per-expert weights ``[N_out, K_in]``, read through ``kernel.cols`` and launched + with ``kernel.dgrad_layout``. A is read per expert in place; the weights are never + concatenated -- the kernel reads expert ``e``'s from ``b_data_ptrs[e]`` (``b_data`` + is one representative, for dtype and shape). The per-tile ``(pid_m, pid_n, expert)`` + descriptor queue is built on the host from ``m_splits``, so there is no + device-to-host sync. Operands arrive already quantized. + """ + if layout == "NN": + b_read, layout = kernel.cols, kernel.dgrad_layout + else: + b_read = kernel.rows + a_datas, a_scales = zip(*(kernel.rows(t) for t in a_tensors)) + b_datas, b_scales = zip(*(b_read(t) for t in b_tensors)) + a_data, a_data_ptrs = _row_operand(list(a_datas), m_splits) + device = a_data.device + b_data_ptrs = grouped_utils.expert_ptrs(list(b_datas), device) + a_data = _as_kernel_dtype(a_data, a_tensors[0]) + b_data = _as_kernel_dtype(b_datas[0], b_tensors[0]) + + # TN: B is [N, K]; NN: B is [K, N]. + n = b_data.size(0) if layout == "TN" else b_data.size(-1) + k = a_data.size(-1) + m_total = sum(int(m) for m in m_splits) + D = _validate_or_allocate_output( + D, + shape=torch.Size((m_total, n)), + dtype=D.dtype, + device=device, + backend_name=f"{kernel.name} grouped {layout}", + ) + + # Tiles cover each expert's rows of the packed output; expert token counts must be + # BLOCK_M multiples (the queue builder raises otherwise), so no tile spans two + # experts. "static" walks the queue with a grid stride; "worksteal" claims tiles + # atomically to balance uneven experts (NVTE_FLYDSL_GROUPED_SCHED overrides). + work_queue, num_tiles = grouped_utils.build_fwd_dgrad_work_queue( + m_splits, n // kernel.block_n, kernel.block_m, device + ) + sched = grouped_utils.select_grouped_sched(num_tiles, device) + kernel.launch( + a_data, + kernel.scale_arg(a_scales, device, None), + b_data, + kernel.scale_arg(b_scales, device, None), + D.view(m_total, n), + layout=layout, + m=m_total, + n=n, + k=k, + group=True, + num_groups=len(m_splits), + work_queue=work_queue, + work_counter=grouped_utils.work_counter(sched, device), + sched=sched, + num_work=num_tiles, + b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, + a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, + ) + + +def _wgrad_operands(dy_b, x_b, unpadded, pad, device): + """One wgrad bucket's (dy, dy_ptrs, x, x_ptrs), shared by every dtype. + + When no expert in the bucket needs padding (``unpadded``), both operands are read + in place from their per-expert tensors (``grouped_utils.stacked_operand``), so the kernel + sees the packed layout without a host-side concat. Otherwise ``pad`` builds the + padded packed operand and the pointers are ``None``. + """ + if unpadded: + dy = grouped_utils.stacked_operand(dy_b, device) + x = grouped_utils.stacked_operand(x_b, device) + if dy is not None and x is not None: + return dy[0], dy[1], x[0], x[1] + return pad(dy_b), None, pad(x_b), None + + +def _wgrad_max_buckets(): + """Launch-count cap for the K-bucketed wgrad (env override, default uncapped).""" + v = os.getenv("NVTE_FLYDSL_WGRAD_MAX_BUCKETS", "") + if v.strip() in ("", "0"): + return None + return int(v) + + +def _pad_tokens(tensors, k_pad, token_axis): + """Zero-pad each expert's tokens to ``k_pad`` and stack the bucket's experts.""" + if token_axis == 0: + return grouped_utils.uniform_pad_concat(tensors, k_pad) + padded = [] + for t in tensors: + # [feature, m_e]: tokens are the trailing axis. + if t.size(1) == k_pad: + padded.append(t.contiguous()) + else: + out = t.new_zeros((t.size(0), k_pad)) + out[:, : t.size(1)].copy_(t) + padded.append(out) + return torch.cat(padded, dim=0) + + +def _run_grouped_wgrad(kernel, dy_tensors, x_tensors, dW, *, m_splits): + """Grouped wgrad: ``dW_e = dY_e^T @ X_e`` per expert, via K-bucketed launches. + + Both operands are read through ``kernel.cols``: ``[m_e, feature]`` for MXFP8 and + half precision (an NT launch), or the materialized ``[feature, m_e]`` transposes for + tensor-wise FP8 (a TN launch). ``dW`` is the packed ``[G*N_out, K_in]`` output. + + The kernel's contraction K (the per-expert token count) is baked in at compile + time, so experts are grouped into **K-buckets** by their own padded token count + (``bucketize_wgrad_tokens``): each expert rounds up to its own BLOCK_M multiple + (floor 512), and experts sharing a rounded K share one launch. A single large + expert therefore no longer forces every small expert to contract over its token + count (which wasted work and overflowed the int32 launch-arg limit under skew). + Balanced token counts collapse to a single bucket. + + Per bucket, the experts' operands are read in place when none needs padding, + otherwise zero-padded on the token axis to the bucket's ``k_pad`` and stacked + (bucket-local, so the kernel's ``expert*K`` addressing uses the bucket-local + index). Output goes straight to ``dW`` when the bucket is every expert in order, + else to a bucket-local buffer copied into each expert's global rows. The + hand-scheduled 4-wave MFMA pipeline is untouched -- K stays compile-time per launch. + """ + dy_datas, dy_scales = zip(*(kernel.cols(t) for t in dy_tensors)) + x_datas, x_scales = zip(*(kernel.cols(t) for t in x_tensors)) + axis = kernel.wgrad_token_axis + n_out = dy_datas[0].size(1 - axis) + k_in = x_datas[0].size(1 - axis) + num_groups = len(m_splits) + device = dW.device + dW = _validate_or_allocate_output( + dW, + shape=torch.Size((num_groups * n_out, k_in)), + dtype=dW.dtype, + device=device, + backend_name=f"{kernel.name} grouped wgrad", + ) + + # One (pid_m, pid_n, expert) descriptor per tile of each expert's [N_out, K_in] + # weight gradient; the kernel offsets the output row by expert*N_out. + num_pid_m = n_out // kernel.block_m + num_pid_n = k_in // kernel.block_n + buckets = grouped_utils.bucketize_wgrad_tokens( + m_splits, kernel.block_m, max_buckets=_wgrad_max_buckets() + ) + buckets = grouped_utils.drop_empty_experts(buckets, m_splits, dW, n_out) + + for k_pad, ids in buckets: + bsize = len(ids) + dy_b = [dy_datas[e] for e in ids] + x_b = [x_datas[e] for e in ids] + dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( + dy_b, + x_b, + all(t.size(axis) == k_pad for t in dy_b + x_b), + lambda ts: _pad_tokens(ts, k_pad, axis), + device, + ) + dy_data = _as_kernel_dtype(dy_data, dy_tensors[0]) + x_data = _as_kernel_dtype(x_data, x_tensors[0]) + # Scales are read in place when the operands are; otherwise padded alongside them. + scale_k_pad = None if dy_ptrs is not None else k_pad + dy_scale = kernel.scale_arg([dy_scales[e] for e in ids], device, scale_k_pad) + x_scale = kernel.scale_arg([x_scales[e] for e in ids], device, scale_k_pad) + + # When the bucket holds all experts in order (the balanced case), bucket-local + # and global expert indices match, so the kernel writes dW directly. + direct = ids == list(range(num_groups)) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=dW.dtype, device=device) + + work_queue, num_tiles = grouped_utils.build_wgrad_work_queue( + bsize, num_pid_m, num_pid_n, device + ) + sched = grouped_utils.select_grouped_sched(num_tiles, device) + if axis == 0: + a_shape, b_shape = (bsize * k_pad, n_out), (bsize * k_pad, k_in) + else: # transposed operands stack the bucket's experts on the feature axis + a_shape, b_shape = (bsize * n_out, k_pad), (bsize * k_in, k_pad) + kernel.launch( + dy_data, + dy_scale, + x_data, + x_scale, + dW_b, + layout="NT", + m=n_out, + n=k_in, + k=k_pad, + group=False, + num_groups=bsize, + work_queue=work_queue, + work_counter=grouped_utils.work_counter(sched, device), + sched=sched, + wgrad=True, + num_work=num_tiles, + a_data_ptrs=dy_ptrs, + b_data_ptrs=x_ptrs, + a_shape=a_shape if dy_ptrs is not None else None, + b_shape=b_shape if x_ptrs is not None else None, + ) + + if not direct: + # Scatter each bucket-local expert block to its global dW rows. + for j, e in enumerate(ids): + dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) + + +def _packed_wgrad_view(wgrad_list, n_out, k_in): + """Return a ``[G*n_out, k_in]`` view over ``wgrad_list`` if it is one packed run. + + The non-fused grouped-wgrad destinations are adjacent ``[n_out, k_in]`` slices of + one contiguous ``[G, n_out, k_in]`` buffer, so the kernel can write its packed + output directly into that storage, eliding the per-expert copy-out. Return ``None`` + (the caller allocates a scratch buffer and scatters) whenever the destinations are + not a single contiguous, row-contiguous, correctly-shaped run. + """ + first = wgrad_list[0] + elem = first.element_size() + expected_ptr = first.data_ptr() + for w in wgrad_list: + if ( + w.dtype != first.dtype + or tuple(w.shape) != (n_out, k_in) + or not w.is_contiguous() + or w.data_ptr() != expected_ptr + ): + return None + expected_ptr += w.numel() * elem + # One contiguous run anchored at first.data_ptr(); span all rows over its storage. + return torch.as_strided(first, (len(wgrad_list) * n_out, k_in), (k_in, 1)) + + +def _grouped_gemm_flydsl(A, B, out, *, layout, m_splits, single_output): + """Run ``general_grouped_gemm``'s plain GEMM on FlyDSL, or raise if unsupported. + + TE's grouped layouts map onto the three phases: TN forward ``out = B_e @ A_e^T`` + (A = weights, B = activations), NN dgrad ``out = B_e @ A_e`` (A = weights, + B = grad_output), NT wgrad ``out[e] = B_e^T @ A_e`` (A = activations, + B = grad_output). Raises ``FlyDSLUnsupportedError`` for any operand or config the + kernels cannot serve. + """ + if layout not in ("TN", "NN", "NT") or m_splits is None: + raise FlyDSLUnsupportedError(f"FlyDSL grouped GEMM does not support layout {layout}") + kernel = _grouped_kernel_for(list(A) + list(B)) + if layout == "NT": + if single_output: + raise FlyDSLUnsupportedError("FlyDSL grouped wgrad writes per-expert outputs") + k_in = kernel.cols(A[0])[0].size(1 - kernel.wgrad_token_axis) + n_out = out[0].size(0) + # Write straight into the packed wgrad storage when ``out`` is one contiguous + # run; otherwise launch into scratch and scatter. + dW = _packed_wgrad_view(out, n_out, k_in) + scatter = dW is None + if scatter: + dW = torch.empty(len(out) * n_out, k_in, dtype=out[0].dtype, device=out[0].device) + _run_grouped_wgrad(kernel, B, A, dW, m_splits=m_splits) + if scatter: + for g, w in enumerate(out): + w.copy_(dW[g * n_out : (g + 1) * n_out]) + return + if not single_output: + raise FlyDSLUnsupportedError("FlyDSL grouped fwd/dgrad writes one packed output") + _run_grouped(kernel, B, A, out[0], m_splits=m_splits, layout=layout) + + +def _warn_fallback_enabled() -> bool: + return os.environ.get("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "0").lower() not in ( + "", + "0", + "false", + "no", + "off", + ) + + +def general_grouped_gemm_flydsl( + A, + B, + out, + quantization_params, + out_dtype, + layout="TN", + m_splits=None, + gelu=False, + grad=False, + accumulate=False, + bias=None, + use_bias=False, + use_split_accumulator=False, + D_dtype=None, + single_output=False, + **kwargs, +): + """Drop-in replacement for ``general_grouped_gemm`` using the FlyDSL grouped kernels. + + Serves the plain grouped GEMM over BF16/FP16, tensor-wise FP8, and MXFP8 operands + (gfx950) in all three layouts. Anything else -- bias, GELU, accumulation, an + explicit output dtype, another quantized storage, or a shape the kernel tiling + cannot cover -- falls back to ``general_grouped_gemm``, mirroring the dense FlyDSL + fallback in ``cpp_extensions.gemm``. ``NVTE_FLYDSL_GEMM_WARN_FALLBACK=1`` warns on + fallback. Returns ``general_grouped_gemm``'s ``(out, bias, gelu_input)`` triple. + """ + from ...cpp_extensions.gemm import _empty_tensor, general_grouped_gemm + + try: + if gelu or accumulate or use_bias or D_dtype is not None: + raise FlyDSLUnsupportedError( + "FlyDSL grouped GEMM does not support gelu, accumulate, bias, or D_dtype" + ) + _grouped_gemm_flydsl( + A, B, out, layout=layout, m_splits=m_splits, single_output=single_output + ) + except FlyDSLUnsupportedError as exc: + if _warn_fallback_enabled(): + warnings.warn( + "[FLYDSL WARNING]: FlyDSL grouped GEMM does not support this configuration " + f"(layout={layout}); falling back to the default backend. Reason: {exc}", + UserWarning, + stacklevel=2, + ) + return general_grouped_gemm( + A, + B, + out, + quantization_params, + out_dtype, + layout=layout, + m_splits=m_splits, + gelu=gelu, + grad=grad, + accumulate=accumulate, + bias=bias, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + D_dtype=D_dtype, + single_output=single_output, + ) + empty_tensors = [_empty_tensor()] * len(A) + return out, empty_tensors, empty_tensors + + def _select_fp8_storage_for_layout(A, transa, B, transb): """Select existing TE FP8 backings and normalize to one core contract. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py new file mode 100644 index 000000000..11bc8694e --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -0,0 +1,516 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Shared utilities for the FlyDSL persistent work-stealing grouped GEMM. + +The grouped (MoE) GEMM path is dtype-independent: one packed kernel launch spans +all local experts, a fixed num-CU persistent grid claims tiles from a precomputed +descriptor queue, and the per-expert offset is folded into operand/scale/output +addressing. MXFP8, BF16/FP16, and tensor-wise FP8 all share this machinery. + +This module holds the reusable pieces that can live outside the ``@flyc.kernel`` +body. The persistent loop itself -- the ``if tx==0`` atomic claim and the ``while`` +drain loops -- is NOT here: the FlyDSL compiler AST-rewrites only the decorated +kernel's own source, so that runtime-SSA control flow MUST stay inline in each +``kernel_gemm``. What's shared instead: + +* **Device-side IR emitters** (``global_i32_ptr``, ``make_claim_broadcast``, + ``read_desc_field``) -- control-flow-free functions that + append IR at the current insertion point, so they emit identically whether + defined inline or imported. +* **Host-side builders** (descriptor-queue construction, schedule selection, + work-counter allocation, uniform token-axis padding) -- plain Python/torch, + identical across dtypes; they differ only in whether scales are also padded. +""" + +import functools +import os + +import numpy as np +import torch + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm +from flydsl.expr.typing import T + +from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import get_num_cus, placeholder_arg + + +# --------------------------------------------------------------------------- +# Device-side IR emitters (called while tracing a kernel; pure, no control flow) +# --------------------------------------------------------------------------- + + +def global_i32_ptr(tensor, index): + """Build an address-space-1 ``!llvm.ptr`` to ``tensor[index]`` (int32). + + ``llvm.atomicrmw`` requires an ``!llvm.ptr`` (not an ``!fly.ptr``); this + recovers the raw base pointer and offsets it by ``index`` int32 elements. + """ + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + base = _fly.extract_aligned_pointer_as_index(ptr_ty, tensor) + base = llvm.PtrToIntOp(T.i64, base).result + byte_off = fx.as_ir_value(fx.Int64(fx.Int64(index) * fx.Int64(4))) + addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result + return llvm.IntToPtrOp(ptr_ty, addr).result + + +def read_ptr_at(ptr_array, index): + """Load the int64 base address stored at ``ptr_array[index]``. + + ``ptr_array`` is a device ``int64[G]`` of per-expert operand base addresses + (each the ``data_ptr()`` of one expert's independently-allocated weight slab). + Recovers the array base, offsets by ``index`` int64 elements (8 bytes), and + loads the stored address as an i64 IR value -- the per-expert operand base. + """ + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + base = _fly.extract_aligned_pointer_as_index(ptr_ty, ptr_array) + base = llvm.PtrToIntOp(T.i64, base).result + byte_off = fx.as_ir_value(fx.Int64(fx.Int64(index) * fx.Int64(8))) + addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result + elt_ptr = llvm.IntToPtrOp(ptr_ty, addr).result + return llvm.LoadOp(T.i64, elt_ptr).result + + +def load_i32_at(tensor, index): + """Load ``tensor[index]`` (int32) with a plain per-lane global load.""" + return llvm.LoadOp(T.i32, global_i32_ptr(tensor, index)).result + + +def make_uniform_i64(value): + """Broadcast lane 0's i64 to the whole wave (``readfirstlane``). + + A buffer descriptor lives in SGPRs. An address loaded by :func:`read_ptr_at` is a + per-lane value, so a descriptor rebased on it can make the compiler guard each + buffer access for divergence. Use this when every lane loaded the same address + (e.g. indexed by the tile's expert), so lane 0's value is every lane's value. + """ + return fx.rocdl.readfirstlane(T.i64, value) + + +def load_f32_at_ptr(ptr_array, index): + """Load the fp32 scalar that ``ptr_array[index]`` points to. + + A plain global load needs no wave-uniform address, so unlike the buffer-descriptor + bases this skips :func:`make_uniform_i64`: nothing waits on either load until the + value is used, which keeps a per-tile scale fetch off the tile's critical path. + """ + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + addr = read_ptr_at(ptr_array, index) + return llvm.LoadOp(T.f32, llvm.IntToPtrOp(ptr_ty, addr).result).result + + +def rebase_buffer_view(base_i64, elem_ir_type, layout): + """Rebuild a buffer-tensor logical-divide view at a runtime base address. + + Reuses an EXISTING static ``layout`` (compile-time cosize) but swaps the base + pointer to the runtime ``base_i64`` address. This is how a per-expert operand + slice is addressed without constructing a fresh runtime-sized layout + (``fx.make_layout`` requires a static extent): the packed operand's layout is + reused and only the base moves. ``max_size=True`` lets the buffer descriptor + span the per-expert extent; the loader's dynamic ``soffset`` walks within it. + """ + buf_ptr_ty = fx.PointerType.get(elem_ty=elem_ir_type, address_space=1) + base_ptr = fx.inttoptr(buf_ptr_ty, fx.Int64(base_i64)) + view = fx.make_view(base_ptr, layout) + buf = fx.rocdl.make_buffer_tensor(fx.Tensor(view), max_size=True) + return fx.logical_divide(buf, fx.make_layout(1, 1)) + + +def make_claim_broadcast(claim_slot_array): + """Recast a one-int32 LDS scratch array into a claim iterator + 1-elem view. + + Thread 0 writes the atomically-claimed slot via ``claim_iter``; after a + workgroup barrier every thread reads the same value through ``claim_view``. + """ + claim_iter = fx.recast_iter(fx.Int32, claim_slot_array.ptr) + claim_view = fx.make_view(claim_iter, fx.make_layout(1, 1)) + return claim_iter, claim_view + + +def read_desc_field(wq_div, desc_ld_atom, slot, field): + """Read one int32 of the packed ``(pid_m, pid_n, expert)`` descriptor. + + ``wq_div`` is the ``logical_divide`` view of the work-queue buffer tensor and + ``desc_ld_atom`` an int32 ``BufferCopy32b`` copy atom. + """ + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(desc_ld_atom, fx.slice(wq_div, (None, slot * fx.Int32(3) + field)), reg) + return fx.memref_load_vec(reg)[0] + + +# --------------------------------------------------------------------------- +# Host-side builders (plain torch; dtype-independent) +# --------------------------------------------------------------------------- + +# Above this many persistent-loop rounds, dynamic atomic work-stealing pays for +# itself; below it a single grid-stride sweep already balances the tiles, so the +# atomic-free static schedule is preferred. Matches the HK reference's SCHED_ROUNDS. +_GROUPED_SCHED_ROUNDS = 4 + + +def select_grouped_sched(num_tiles: int, device) -> str: + """Pick the grouped-GEMM schedule: 'worksteal' (atomic) or 'static' (grid-stride). + + Honors the ``NVTE_FLYDSL_GROUPED_SCHED`` override ('auto', 'static', + 'worksteal') for testing; otherwise applies the round-count heuristic against + the CU grid. + """ + override = os.getenv("NVTE_FLYDSL_GROUPED_SCHED", "auto").lower() + if override in ("static", "worksteal"): + return override + if override not in ("auto", ""): + raise ValueError( + f"NVTE_FLYDSL_GROUPED_SCHED must be auto|static|worksteal, got {override!r}" + ) + return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * get_num_cus(device) else "worksteal" + + +# Per-call metadata caches. Keys include the current stream: a cached device tensor is +# filled by a non-blocking copy on the stream that first built it, so only launches on +# that stream are ordered after the copy. Pointer arrays and block tables are tiny +# (<= a few KB each), so their cache is large enough to hold a training step's worth of +# rotating per-expert buffers (each step adds several arrays); the work queues can be +# much larger, so their cache stays small. +_SMALL_CACHE_SIZE = 512 +_QUEUE_CACHE_SIZE = 64 + + +def current_stream_handle(device=None): + """Raw handle of the current stream on ``device`` (default: the current device). + + FlyDSL launchers accept an int stream, and this is ~20x cheaper than building a + torch ``Stream`` object per launch. + """ + index = None if device is None else device.index + if index is None: + index = torch.cuda.current_device() + return torch._C._cuda_getCurrentRawStream(index) + + +def upload_to_device(host, device): + """Copy a host CPU tensor to ``device`` without blocking the host. + + ``torch.tensor(..., device="cuda")`` waits for all queued GPU work before copying, + which stalls the host on every call. Staging through pinned memory and copying + ``non_blocking`` does not; the caching host allocator keeps the pinned block alive + until the copy has run. The staging block comes from ``torch.empty(pin_memory=True)`` + (the caching host allocator) rather than ``Tensor.pin_memory()``, which is ~3x slower + per call while GPU work is queued. + """ + staging = torch.empty(host.shape, dtype=host.dtype, pin_memory=True) + staging.copy_(host) + return staging.to(device, non_blocking=True) + + +@functools.lru_cache(maxsize=_QUEUE_CACHE_SIZE) +def _fwd_dgrad_queue_cached(m_splits, num_pid_n, block_m, device, stream): + """Flat ``(pid_m, pid_n, expert)`` work queue for a fwd/dgrad shape. + + Every expert's token count is a ``block_m`` multiple, so the expert row blocks tile + ``[0, M_total)`` in order: the queue is each row block repeated over the output + columns, tagged with its expert. Built with numpy from the host ``m_splits`` (no + device->host copy of ``group_offs``) and uploaded without blocking the host. + """ + blocks = np.asarray(m_splits, dtype=np.int64) // block_m + num_blocks = int(blocks.sum()) + queue = np.empty((num_blocks, num_pid_n, 3), dtype=np.int32) + queue[:, :, 0] = np.arange(num_blocks, dtype=np.int32)[:, None] + queue[:, :, 1] = np.arange(num_pid_n, dtype=np.int32)[None, :] + queue[:, :, 2] = np.repeat(np.arange(len(m_splits), dtype=np.int32), blocks)[:, None] + work_queue = upload_to_device(torch.from_numpy(queue.reshape(-1)), device) + return work_queue, num_blocks * num_pid_n + + +def build_fwd_dgrad_work_queue(m_splits, num_pid_n, block_m, device): + """Return ``(work_queue, num_tiles)`` for a fwd/dgrad launch. + + Raises ``FlyDSLUnsupportedError`` when an expert's token count is not a ``block_m`` + multiple: the kernel tiles each expert in whole ``block_m`` row blocks, so a + misaligned split would silently mix experts within a tile. ``num_tiles`` is the + exact queue length, which the kernel uses as its drain bound. ``m_splits`` is a + Python list/tuple of per-expert token counts -- no GPU->CPU sync. + """ + m_splits = tuple(int(m) for m in m_splits) + for g, m in enumerate(m_splits): + if m % block_m != 0: + raise FlyDSLUnsupportedError( + f"m_splits[{g}]={m} is not a multiple of BLOCK_M={block_m}" + ) + return _fwd_dgrad_queue_cached( + m_splits, int(num_pid_n), int(block_m), device, current_stream_handle(device) + ) + + +@functools.lru_cache(maxsize=_QUEUE_CACHE_SIZE) +def _wgrad_queue_cached(num_groups, num_pid_m, num_pid_n, device, stream): + queue = np.empty((num_groups, num_pid_m, num_pid_n, 3), dtype=np.int32) + queue[..., 0] = np.arange(num_pid_m, dtype=np.int32)[None, :, None] + queue[..., 1] = np.arange(num_pid_n, dtype=np.int32)[None, None, :] + queue[..., 2] = np.arange(num_groups, dtype=np.int32)[:, None, None] + work_queue = upload_to_device(torch.from_numpy(queue.reshape(-1)), device) + return work_queue, num_groups * num_pid_m * num_pid_n + + +def build_wgrad_work_queue(num_groups, num_pid_m, num_pid_n, device): + """Return ``(work_queue, num_tiles)`` for a wgrad launch. + + The wgrad descriptor queue depends only on ``(num_groups, num_pid_m, num_pid_n)`` + (per-expert tile geometry), so it is shape-cached and reused across calls; + ``num_tiles`` is the kernel's drain bound. + """ + return _wgrad_queue_cached( + int(num_groups), int(num_pid_m), int(num_pid_n), device, current_stream_handle(device) + ) + + +def work_counter(sched, device): + """The atomic tile counter a grouped launch needs for schedule ``sched``. + + Work-stealing claims tiles from a freshly zeroed counter. The static schedule never + touches it, so it gets a shared placeholder instead of a per-launch zero fill. + """ + if sched == "worksteal": + return torch.zeros(1, dtype=torch.int32, device=device) + return placeholder_arg(device, torch.int32) + + +class _SpanCAI: + """Minimal ``__cuda_array_interface__`` holder for a raw uint8 device span. + + Wraps ``total_bytes`` at ``data_ptr`` as a 1-D uint8 CAI so ``torch.as_tensor`` + can build a zero-copy view over a contiguous multi-tensor run (see + :func:`concat_or_view_rows`). Hoisted to module scope so the class object is + built once at import, not reconstructed on every call. + """ + + __slots__ = ("__cuda_array_interface__",) + + def __init__(self, total_bytes, data_ptr): + self.__cuda_array_interface__ = { + "shape": (total_bytes,), + "typestr": "|u1", + "data": (data_ptr, False), + "strides": None, + "version": 3, + } + + +def concat_or_view_rows(tensors): + """Concatenate per-expert tensors on dim 0, skipping the copy when possible. + + ``tex.split_quantize`` bulk-allocates its per-expert outputs into one base + buffer (data blocks packed back-to-back, then scale blocks back-to-back), so + the per-expert rowwise payloads are frequently already physically adjacent -- + the ``[M_total, cols]`` concatenation the grouped kernel wants already exists + in memory. When that holds, return a zero-copy spanning view over the shared + region instead of materializing a redundant ``torch.cat`` copy (the activation + and activation-scale cats on the MXFP8 forward/dgrad hot path). + + Adjacency is not guaranteed (independently-allocated operands -- e.g. weights + from separate ``nn.Parameter`` storages -- are scattered), so fall back to + ``torch.cat`` whenever the inputs are not a single contiguous run. The view and + the cat produce bitwise-identical operands; only the copy is elided. + """ + if len(tensors) == 1: + return tensors[0] + + first = tensors[0] + if not first.is_cuda or first.ndim < 1: + return torch.cat(tensors, dim=0) + + # __cuda_array_interface__ is unavailable on grad-tracking tensors. The grouped + # GEMM operand is consumed inside a custom autograd Function (grad handled at the + # module level), so a detached data view is correct; detach once up front so the + # adjacency walk and span wrapper both see plain tensors. + if any(t.requires_grad for t in tensors): + tensors = [t.detach() for t in tensors] + first = tensors[0] + + trailing = first.shape[1:] + elem = first.element_size() + total_rows = 0 + expected_ptr = first.data_ptr() + for t in tensors: + # Adjacency + uniformity: same trailing shape/dtype, row-contiguous, and + # starting exactly where the previous tensor ended. + if ( + t.dtype != first.dtype + or t.shape[1:] != trailing + or not t.is_contiguous() + or t.data_ptr() != expected_ptr + ): + return torch.cat(tensors, dim=0) + total_rows += t.size(0) + expected_ptr += t.numel() * elem + + # One contiguous run: wrap [total_rows*row_bytes] over the shared base pointer as + # raw uint8, then reinterpret back to the operand dtype. Going through bytes keeps + # this dtype-agnostic -- __cuda_array_interface__ has no typestr for bfloat16, so a + # native-dtype CAI would reject it; uint8 is always representable. + # The CAI carries no ownership, so anchor the source tensors on the result to keep + # their bulk buffer alive for the view's lifetime. + row_bytes = 1 + for d in trailing: + row_bytes *= d + row_bytes *= elem + total_bytes = total_rows * row_bytes + + byte_view = torch.as_tensor(_SpanCAI(total_bytes, first.data_ptr()), device=first.device) + view = byte_view.view(first.dtype).view(total_rows, *tuple(trailing)) + view._flydsl_span_anchor = tensors + return view + + +@functools.lru_cache(maxsize=_SMALL_CACHE_SIZE) +def _ptr_array_cached(ptrs, device, stream): + return upload_to_device(torch.tensor(ptrs, dtype=torch.int64), device) + + +def expert_ptrs(tensors, device): + """int64[G] device array of each per-expert tensor's ``data_ptr()``. + + Grouped kernels read per-expert operands (weights, FP8 scale_inv) through this + array instead of a host-side concat. Repeated pointer sets reuse the uploaded array; + a new set costs one non-blocking upload. + """ + return _ptr_array_cached( + tuple(t.data_ptr() for t in tensors), device, current_stream_handle(device) + ) + + +def stacked_operand(tensors, device): + """Address per-expert tensors stacked on dim 0 in place, as one packed operand. + + Returns ``(representative, ptrs)``, or ``None`` when the tensors are not + contiguous with a common trailing shape and dtype. Entry ``g`` of the int64[G] + ``ptrs`` is ``tensors[g].data_ptr()`` minus the byte offset where expert ``g`` + starts in the ``torch.cat(tensors)`` layout, so a kernel keeps its packed + addressing and only swaps the buffer base per tile. ``representative`` is a + non-empty expert tensor standing in for the operand's dtype. + """ + first = tensors[0] + dtype = first.dtype + trailing = first.shape[1:] + ptrs = [] + offset = 0 + for t in tensors: + if t.dtype is not dtype or t.shape[1:] != trailing or not t.is_contiguous(): + return None + ptrs.append(t.data_ptr() - offset) + offset += t.nbytes + representative = next((t for t in tensors if t.numel() > 0), first) + return representative, _ptr_array_cached(tuple(ptrs), device, current_stream_handle(device)) + + +@functools.lru_cache(maxsize=_SMALL_CACHE_SIZE) +def _block_expert_table_cached(blocks_per_expert, device, stream): + table = np.repeat(np.arange(len(blocks_per_expert), dtype=np.int32), blocks_per_expert) + return upload_to_device(torch.from_numpy(table), device) + + +def block_expert_table(extents, block, device): + """int32 table mapping each ``block``-sized slice of a stacked axis to its expert. + + ``extents`` are the per-expert lengths along the stacked axis; each must be a + multiple of ``block``, so no block straddles two experts. + """ + for n in extents: + if int(n) % block != 0: + raise ValueError(f"per-expert extent {int(n)} is not a multiple of {block}") + return _block_expert_table_cached( + tuple(int(n) // block for n in extents), device, current_stream_handle(device) + ) + + +def drop_empty_experts(buckets, m_splits, dW, n_out): + """Zero the ``dW`` rows of zero-token experts and drop them from the wgrad buckets. + + An expert with no tokens has ``dW_e = 0``; launching it anyway contracts a zero + padding that a non-finite scale (FP8 never writes an empty tensor's scale_inv) + would turn into NaN. + """ + for e, m in enumerate(m_splits): + if int(m) == 0: + dW[e * n_out : (e + 1) * n_out].zero_() + buckets = [(k_pad, [e for e in ids if int(m_splits[e]) != 0]) for k_pad, ids in buckets] + return [(k_pad, ids) for k_pad, ids in buckets if ids] + + +def uniform_pad_concat(tensors, rows): + """Zero-pad each tensor's leading dim up to ``rows`` and concat on dim 0. + + Used by the wgrad uniform-pad path to stack per-expert operands (or their + scales) on the contraction/token axis. Padded rows are zero, so they + contribute nothing to the token-axis contraction. + + When no tensor actually needs padding (every leading dim already equals + ``rows`` -- the balanced / no-skew case, where each expert's token count is a + multiple of the bucket's BLOCK_M), there is nothing to zero-fill, so this + degenerates to a plain concat. Delegate that case to ``concat_or_view_rows``, + which returns a zero-copy spanning view when the per-expert payloads are + physically adjacent (the common bulk-allocated layout), eliminating the + ``torch.cat`` copy on the wgrad hot path. + """ + if all(t.size(0) == rows for t in tensors): + return concat_or_view_rows(tensors) + + def _pad_rows(t): + if t.size(0) == rows: + return t.contiguous() + out = t.new_zeros((rows, *t.shape[1:])) + out[: t.size(0)].copy_(t) + return out + + return torch.cat([_pad_rows(t) for t in tensors], dim=0) + + +def bucketize_wgrad_tokens(m_splits, block_m, max_buckets=None): + """Group experts by their per-expert padded contraction length (K-bucketing). + + The grouped wgrad kernel bakes its contraction length ``K`` in at compile time, + so experts that share a padded token count can share one launch. Instead of + padding every expert to the GLOBAL max token count (which wastes work and can + OOM under skew), round each expert's ``m_g`` UP to its own ``block_m`` multiple + (floor 512 = the pipeline's 4-K128-tile minimum), then group experts by that + rounded value. Each distinct value is one bucket -> one fixed-K launch. + + When the number of distinct values exceeds ``max_buckets``, adjacent buckets + are merged UPWARD (experts promoted to the larger K) so the launch count is + capped; promoting only ever adds zero padding, never truncates, so results are + unchanged. Returns a list of ``(k_pad, [global_expert_ids])`` tuples, ordered + by ascending ``k_pad``. + """ + + def _round(m): + return max(((int(m) + block_m - 1) // block_m) * block_m, 512) + + k_of = [_round(m) for m in m_splits] + distinct = sorted(set(k_of)) + + if max_buckets is not None and len(distinct) > max_buckets: + # Merge the smallest-K buckets upward onto a coarser ladder until the + # distinct-K count fits. Keep the largest K (one dominant expert) its own + # bucket so it is never dragged smaller (it can't be) nor inflates others. + # Simple policy: collapse the sorted distinct values into `max_buckets` + # contiguous groups, each group promoted to its group-max. + import math + + per = math.ceil(len(distinct) / max_buckets) + ladder = {} + for i in range(0, len(distinct), per): + grp = distinct[i : i + per] + top = grp[-1] + for v in grp: + ladder[v] = top + k_of = [ladder[k] for k in k_of] + distinct = sorted(set(k_of)) + + buckets = [] + for k_pad in distinct: + ids = [g for g, k in enumerate(k_of) if k == k_pad] + buckets.append((k_pad, ids)) + return buckets diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 36ed68d5b..6742324e1 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -7,11 +7,11 @@ FP16 and BF16 share one source-level kernel generator: the algorithm is identical and only the ``v_mfma_f32_16x16x32_{f16,bf16}`` opcode differs. That opcode is selected at compile time via ``mfma_suffix``, so each (dtype, K, -output, layout) combination still compiles to its own cached binary with no +output, layout) combination still compiles to its own cached kernel with no runtime dtype branch. All supported layouts share the same generator while compiling to separate -cached binaries: +cached kernels: TN: A [M,K] normal read, B [N,K] normal read NN: A [M,K] normal read, B [K,N] transpose read @@ -37,8 +37,15 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + get_num_cus, + placeholder_arg, + require_block_tiling, + require_launch_size, +) from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -110,6 +117,10 @@ def _compile_kernel( mfma_suffix: str, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, + act_ptrs: bool = False, ): """Build one compile-time-specialized TN, NN, or NT half-precision kernel. @@ -204,9 +215,9 @@ def _a_leading_dim_bytes(c_m): return c_m * ELEM_BYTES def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): - return k_base * fx.Index(c_m * ELEM_BYTES) + ( - bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) - ) * fx.Index(ELEM_BYTES) + return k_base * fx.Int64(c_m * ELEM_BYTES) + ( + bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2)) + ) * fx.Int64(ELEM_BYTES) def _load_a_half( load_transposed_frag_half, @@ -219,11 +230,11 @@ def _load_a_half( lane_mod_16, ): del load_frag_half_at_byte_base, lane_mod_16 - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) local_m_tile = ( - subtile_m_idx * fx.Index(SUBTILE_M) - + fx.Index(mi * MFMA_M) - - fx.Index(sm * (BLOCK_M // 2)) + subtile_m_idx * fx.Int64(SUBTILE_M) + + fx.Int64(mi * MFMA_M) + - fx.Int64(sm * (BLOCK_M // 2)) ) return load_transposed_frag_half(lds_a[sm], local_m_tile, half) @@ -235,9 +246,9 @@ def _a_leading_dim_bytes(c_m): def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): del c_m - return (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index( + return (bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) def _load_a_half( load_transposed_frag_half, @@ -250,12 +261,12 @@ def _load_a_half( lane_mod_16, ): del load_transposed_frag_half - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) return load_frag_half_at_byte_base( lds_a[sm], - half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half_row * fx.Int64(BLOCK_K * ELEM_BYTES), half, ) @@ -265,9 +276,9 @@ def _b_leading_dim_bytes(c_n): return c_n * ELEM_BYTES def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): - return k_base * fx.Index(c_n * ELEM_BYTES) + ( - by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) - ) * fx.Index(ELEM_BYTES) + return k_base * fx.Int64(c_n * ELEM_BYTES) + ( + by_n_idx + fx.Int64(subtile * (BLOCK_N // 2)) + ) * fx.Int64(ELEM_BYTES) def _load_b_ni( load_transposed_frag, @@ -279,11 +290,11 @@ def _load_b_ni( lane_mod_16, ): del load_normal_b_frag, lane_mod_16 - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) + subtile_n_idx * fx.Int64(SUBTILE_N) + + fx.Int64(ni * MFMA_N) + - fx.Int64(sn * (BLOCK_N // 2)) ) return load_transposed_frag(lds_b[sn], local_n_tile) @@ -295,9 +306,9 @@ def _b_leading_dim_bytes(c_n): def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): del c_n - return (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index( + return (by_n_idx + fx.Int64(subtile * (BLOCK_N // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) def _load_b_ni( load_transposed_frag, @@ -309,8 +320,8 @@ def _load_b_ni( lane_mod_16, ): del load_transposed_frag - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 return load_normal_b_frag(lds_b, b_row_addr, sn) # Resolve global staging maps before FlyDSL captures ``kernel_gemm``. @@ -372,6 +383,9 @@ class SharedStorage: b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + # Persistent grouped path: one-int32 scratch to broadcast the atomically + # claimed queue slot from thread 0 to the whole workgroup. + claim_slot: fx.Array[fx.Int32, 1, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -382,6 +396,12 @@ def kernel_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -397,821 +417,970 @@ def kernel_gemm( a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) tx = gpu.thread_id("x") + tx_i32 = fx.Int32(tx) num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - if const_expr(use_xcd_remap): - pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) - else: - pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) - - bx_m = pid_m * BLOCK_M - by_n = pid_n * BLOCK_N - - # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines - # these values with i32 constants, so Index-typed coordinates would make - # arith.addi receive mixed operand types. - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # Offsets are always bytes. TN uses the original 128-byte XOR - # swizzle. NN/NT stage K-major 16-bit data as two [K64, X64] slices for - # ds_read_b64_tr_b16; the layout choice was resolved before capture. - gl_off_a = _a_global_offsets(lane, wave_id, c_m) - gl_off_b = _b_global_offsets(lane, wave_id, c_n) - - a_g2s = G2SLoader( - a_div, - gl_off_a, - LOAD_PASSES_HALF, - fx.Uint8.ir_type, - wave_id, - ) - b_g2s = G2SLoader( - b_div, - gl_off_b, - LOAD_PASSES_HALF, - fx.Uint8.ir_type, - wave_id, - ) - s2r = S2RLoader(fx.Int32(0), 1) - - layout_lane16 = fx.make_layout((4, 16), (16, 1)) - coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) - lane_div_16 = fx.get(coord_lane16, 0) - lane_mod_16 = fx.get(coord_lane16, 1) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - # Bias is a length-N fp32 vector indexed by the global output-feature - # (N) coordinate and broadcast across the M/token rows. const_expr folds - # this compile-time flag at trace time so the setup (and load_bias) are - # inlined into the kernel scope with no runtime dispatch branch. - if const_expr(has_bias): - gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) - bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) - bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - - def load_bias(col): - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) - return fx.memref_load_vec(reg)[0] - - # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N - # output so the backward pass can recompute the GELU gradient. Same tile - # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. - if const_expr(has_gelu): - gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) - aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) - aux_store_atom = fx.make_copy_atom( + # ``grouped`` gates the shared descriptor queue / work-steal driver (group + # fwd/dgrad or the wgrad specialization); const_expr folds it all out of the + # dense compiled kernel. + grouped = group or wgrad + # Grouped drain bound: the exact number of descriptors the host wrote to + # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. + num_tiles = num_work + + if const_expr(grouped): + wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) + wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + + # Grouped fwd/dgrad: weights arrive as G independently-allocated per-expert + # slabs, so B is addressed from a per-expert base pointer read from + # ``b_data_ptrs`` (byte-addressed, reusing gB's static layout) rather than + # an expert-shifted offset into one packed B. wgrad reads B by pointer only + # with act_ptrs; dense never enters here. + b_ptr_layout = fx.get_layout(gB) + # act_ptrs: the per-expert activation operand (A; and in wgrad also B, via + # b_data_ptrs) stays in its own per-expert allocation instead of a host-side + # concat. The host passes each expert's base pointer pre-shifted by that + # expert's packed offset, so the packed addressing below is unchanged; only + # the buffer base moves per tile. + if const_expr(act_ptrs): + a_ptr_layout = fx.get_layout(gA) + + # LDS broadcast scratch (one int32): thread 0 writes the atomically + # claimed slot, a barrier broadcasts it to the whole workgroup. The + # ``if tx==0`` claim control flow must stay inline (AST-rewritten here). + claim_iter, claim_view = grouped_utils.make_claim_broadcast(lds.claim_slot) + + def _claim_next_slot(): + if tx_i32 == fx.Int32(0): + old = llvm.AtomicRMWOp( + llvm.AtomicBinOp.add, + grouped_utils.global_i32_ptr(work_counter, fx.Int32(0)), + fx.as_ir_value(fx.Int32(1)), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + fx.ptr_store(fx.Vector.from_elements([fx.Int32(old)], fx.Int32), claim_iter) + fx.barrier() + return fx.Vector(claim_view.load())[0] + + def _desc_field(slot, field): + return grouped_utils.read_desc_field(wq_div, desc_ld_atom, slot, field) + + def process_tile(pid_m, pid_n, expert): + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) + + # Per-expert offsets into the packed operands / output. Only wgrad (NT) needs + # them: both operands are packed on the contraction (token) axis + # (a/b_expert_k_off = expert*K) and the output on M (store_m_expert_off). + # Grouped fwd/dgrad read B per expert through b_data_ptrs (see b_div_tile + # below), so they use local coordinates and zero offsets, like dense. + if const_expr(wgrad): + by_n_b_idx = by_n_idx + a_expert_k_off = fx.Int64(expert * fx.Int32(K)) + b_expert_k_off = fx.Int64(expert * fx.Int32(K)) + store_m_expert_off = fx.Int64(expert * fx.Int32(c_m)) + else: + by_n_b_idx = by_n_idx + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) + store_m_expert_off = fx.Int64(0) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # Offsets are always bytes. TN uses the original 128-byte XOR + # swizzle. NN/NT stage K-major 16-bit data as two [K64, X64] slices for + # ds_read_b64_tr_b16; the layout choice was resolved before capture. + gl_off_a = _a_global_offsets(lane, wave_id, c_m) + gl_off_b = _b_global_offsets(lane, wave_id, c_n) + + # Issue both per-expert base-pointer loads before either is made uniform, so a + # tile waits on one memory round trip rather than two. + if const_expr(act_ptrs): + a_base_raw = grouped_utils.read_ptr_at(a_data_ptrs, expert) + if const_expr(group or act_ptrs): + b_base_raw = grouped_utils.read_ptr_at(b_data_ptrs, expert) + if const_expr(act_ptrs): + a_base_i64 = grouped_utils.make_uniform_i64(a_base_raw) + a_div_tile = grouped_utils.rebase_buffer_view( + a_base_i64, fx.Uint8.ir_type, a_ptr_layout + ) + else: + a_div_tile = a_div + a_g2s = G2SLoader( + a_div_tile, + gl_off_a, + LOAD_PASSES_HALF, + fx.Uint8.ir_type, + wave_id, + ) + # Grouped B-data: rebase onto this tile's expert weight (byte buffer); the + # per-expert base already selects the expert, so in-tile addressing is + # local (column by_n_idx, no contraction shift) for both TN and NN. + # wgrad bases are pre-shifted (act_ptrs), so they keep the packed column + # and contraction shift. Dense keeps the shared packed view. + if const_expr(group): + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view( + b_base_i64, fx.Uint8.ir_type, b_ptr_layout + ) + b_data_col = by_n_idx + b_k_off = fx.Int64(0) + elif const_expr(act_ptrs): + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view( + b_base_i64, fx.Uint8.ir_type, b_ptr_layout + ) + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off + else: + b_div_tile = b_div + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off + b_g2s = G2SLoader( + b_div_tile, + gl_off_b, + LOAD_PASSES_HALF, + fx.Uint8.ir_type, + wave_id, + ) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). store_m_expert_off + # (grouped wgrad only) repacks the output row to this expert's block; + # it is 0 for dense and for grouped fwd/dgrad. + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Int64(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), output_fx_dtype, ) - def gelu_tanh(x): - # tanh-approx GELU (matches PyTorch approximate='tanh' and the - # FlyDSL preshuffle reference), expressed through a non-positive - # exponent so exp() cannot overflow: - # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = x * x * x - y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * x * (numerator * (one_f32 / denom)) - - PIN_ACC_BASE = 0 - - def _reg_list(prefix, start, end): - return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) - - def reserve_pinned_accumulators(): - # Reserve a fixed physical AGPR bank for all accumulators. In the - # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, - # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator - # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the - # scaled MFMA accumulation in place and avoids those transfers and spills. - # - # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, - # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. - clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) - llvm.InlineAsmOp( - None, - [], - "", - clobbers, - has_side_effects=True, - ) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy32b() + if output_element_bytes == 4 + else fx.rocdl.BufferCopy16b() + ), + output_fx_dtype, + ) - def zero_pinned_accumulators(): - for ai in range_constexpr(ACCS_PER_WAVE * 4): + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) llvm.InlineAsmOp( None, [], - f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", - f"~{{a{PIN_ACC_BASE + ai}}}", + "", + clobbers, has_side_effects=True, ) - def _inline_asm_i32(asm_string, constraints, operands=None): - op = llvm.InlineAsmOp( - T.i32, - operands or [], - asm_string, - constraints, - has_side_effects=True, - ) - return _one_i32_result(op) - - def _one_i32_result(op): - # Accept the result attribute names exposed by the supported MLIR Python bindings. - return getattr(op, "result", getattr(op, "res", op.results[0])) - - def read_pinned_accumulator(acc_idx): - acc_pin = PIN_ACC_BASE + acc_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def read_physical_accumulator_slot(slot_idx): - acc_pin = PIN_ACC_BASE + slot_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def hot_loop_scheduler_q_refill_2n(): - # Eight refill VMEM operations overlap four independent 8-MFMA - # groups (two K32 slices x two N-halves). - for _ in range_constexpr(4): - rocdl.sched_vmem(2) - rocdl.sched_mfma(8) - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q0_refill_a1_2n(): - # Eight refill VMEM operations and eight distributed A-bottom LDS - # reads overlap four independent 8-MFMA K32 groups. - for _ in range_constexpr(4): - rocdl.sched_vmem(2) - rocdl.sched_dsrd(Q0_SCHED_DSRD) - rocdl.sched_mfma(8) - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q_prefetch_4n(): - # Eight two-read prefetch groups overlap four complete-quadrant - # 16-MFMA groups (two K32 slices for each of Q2 and Q3). - for _ in range_constexpr(4): - rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) - rocdl.sched_mfma(16) - rocdl.sched_barrier(0) - - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one - # 128x64 half-page (16 KiB). Each half has its own LDS base. - global_base = _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx) - a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = _b_global_base_bytes(k_base, subtile, c_n, by_n_idx) - b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_a_subtile(k_base, subtile, lds_a): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - - def stage_b_subtile(k_base, subtile, lds_b): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) - - def load_frag_half_at_byte_base(lds_page, row_byte_base, half): - # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. - # Keeping the halves separate allows steady-state Q0 to schedule one - # A-bottom ds_read_b128 in each refill/MFMA chunk. - k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 - return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) - - def pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_frag_at_byte_base(lds_page, row_byte_base): - # Default complete-fragment path used outside the dedicated Q0 schedule. - x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) - x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) - return pack_frag_halves(x0, x1) - - def load_b_frag(lds_b, local_row, half): - # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base( - lds_b[half], - half_row * fx.Index(BLOCK_K * ELEM_BYTES), - ) - - def load_transposed_frag_half(lds_page, local_x_tile, half): - # Half-precision uses v_mfma_f32_16x16x32_{f16,bf16}, not the MXFP8 K128 - # instruction. A 128-X half-page is therefore two independent - # swizzled [K64, X64] half-precision slices. One ds_read_b64_tr_b16 returns - # four half-precision values/lane; two reads form one K32 MFMA fragment. - local_x_i32 = fx.Int32(local_x_tile) - slice_idx = local_x_i32 // fx.Int32(64) - x_in_slice = local_x_i32 % fx.Int32(64) - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - - source_k = lane_div16_i32 * fx.Int32(8) + lane_in16_i32 // fx.Int32(4) - source_x_byte = x_in_slice * fx.Int32(ELEM_BYTES) + ( - lane_in16_i32 % fx.Int32(4) - ) * fx.Int32(8) - - physical_k, physical_x = swizzle_128(source_k, source_x_byte) - slice_base = slice_idx * fx.Int32(64 * 128) - base = slice_base + physical_k * fx.Int32(128) + physical_x - other = base ^ fx.Int32(0x220) - immediate_offset = 0 if half == 0 else 0x1000 - return s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=immediate_offset, - ) + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) - def load_transposed_frag(lds_page, local_x_tile): - x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) - x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) - return pack_frag_halves(x0, x1) + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K32 slices x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(Q0_SCHED_DSRD) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K32 slices for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) + rocdl.sched_mfma(16) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x64 half-page (16 KiB). Each half has its own LDS base. + # a_expert_k_off shifts the contraction base to this tile's expert A + # block for grouped wgrad (dY packed on the token axis); 0 otherwise. + global_base = _a_global_base_bytes(k_base + a_expert_k_off, subtile, c_m, bx_m_idx) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # Grouped fwd/dgrad select the expert via the rebased per-expert B view + # (b_g2s) with LOCAL coords: TN a local column (b_data_col), NN a zero + # contraction shift (b_k_off). Dense/wgrad use the shared view with the + # packed column by_n_b_idx and expert contraction shift b_k_off. + global_base = _b_global_base_bytes(k_base + b_k_off, subtile, c_n, b_data_col) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Int64(BLOCK_K * ELEM_BYTES), + ) - def _acc_idx(subtile_id, mi, ni): - return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni - ) + def load_transposed_frag_half(lds_page, local_x_tile, half): + # Half-precision uses v_mfma_f32_16x16x32_{f16,bf16}, not the MXFP8 K128 + # instruction. A 128-X half-page is therefore two independent + # swizzled [K64, X64] half-precision slices. One ds_read_b64_tr_b16 returns + # four half-precision values/lane; two reads form one K32 MFMA fragment. + local_x_i32 = fx.Int32(local_x_tile) + slice_idx = local_x_i32 // fx.Int32(64) + x_in_slice = local_x_i32 % fx.Int32(64) + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + + source_k = lane_div16_i32 * fx.Int32(8) + lane_in16_i32 // fx.Int32(4) + source_x_byte = x_in_slice * fx.Int32(ELEM_BYTES) + ( + lane_in16_i32 % fx.Int32(4) + ) * fx.Int32(8) + + physical_k, physical_x = swizzle_128(source_k, source_x_byte) + slice_base = slice_idx * fx.Int32(64 * 128) + base = slice_base + physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x220) + immediate_offset = 0 if half == 0 else 0x1000 + return s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) - def _k32_frag(full_frag, k32): - # A/B 16x64 half-precision wave fragments are i32x8. Each K32 MFMA - # consumes one contiguous i32x4 slice (eight half-precision values/lane). - lo = k32 * 4 - v = Vec(full_frag) - return Vec.from_elements( - [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], - fx.Int32, - ) + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) + return pack_frag_halves(x0, x1) - def _pinned_mfma_once(acc_idx, a_k32, b_k32): - acc_pin = PIN_ACC_BASE + acc_idx * 4 - llvm.InlineAsmOp( - None, - [arith._to_raw(a_k32), arith._to_raw(b_k32)], - f"v_mfma_f32_16x16x32_{mfma_suffix} " - f"a[{acc_pin}:{acc_pin + 3}], " - "$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}]", - f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", - has_side_effects=True, - ) + def _acc_idx(subtile_id, mi, ni): + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + mi * MFMA_N_PER_SUBTILE + + ni + ) - def pinned_mfma(acc_idx, a_frag, b_frag): - """Accumulate one logical 16x16x64 half-precision product into pinned AGPRs.""" - for k32 in range_constexpr(2): - _pinned_mfma_once( - acc_idx, - _k32_frag(a_frag, k32), - _k32_frag(b_frag, k32), + def _k32_frag(full_frag, k32): + # A/B 16x64 half-precision wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight half-precision values/lane). + lo = k32 * 4 + v = Vec(full_frag) + return Vec.from_elements( + [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], + fx.Int32, ) - def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): - # The final logical K64 update is two in-place K32 MFMAs. - assert dst_slot == old_acc_idx - pinned_mfma(old_acc_idx, a_frag, b_frag) - - def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - pinned_mfma(acc_base + 2, a_frag, b2) - pinned_mfma(acc_base + 3, a_frag, b3) - - def mfma_2n(acc_base, a_frag, b0, b1): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - - def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): - """Issue one K32 slice for a 4x2 accumulator slab.""" - a_frags = (a0, a1, a2, a3) - b_frags = (b0, b1) - for mi in range_constexpr(4): - a_k32 = _k32_frag(a_frags[mi], k32) - for nj in range_constexpr(2): - _pinned_mfma_once( - _acc_idx(subtile_id, mi, n_base + nj), - a_k32, - _k32_frag(b_frags[nj], k32), - ) + def _pinned_mfma_once(acc_idx, a_k32, b_k32): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k32), arith._to_raw(b_k32)], + f"v_mfma_f32_16x16x32_{mfma_suffix} " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]", + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", + has_side_effects=True, + ) - def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): - """Issue one K32 slice for a complete 4x4 quadrant.""" - a_frags = (a0, a1, a2, a3) - b_frags = (b0, b1, b2, b3) - for mi in range_constexpr(4): - a_k32 = _k32_frag(a_frags[mi], k32) - for ni in range_constexpr(4): + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 half-precision product into pinned AGPRs.""" + for k32 in range_constexpr(2): _pinned_mfma_once( - _acc_idx(subtile_id, mi, ni), - a_k32, - _k32_frag(b_frags[ni], k32), + acc_idx, + _k32_frag(a_frag, k32), + _k32_frag(b_frag, k32), ) - def store_acc_vector_for_logical_idx(logical_acc_idx, acc): - subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - - # Bias depends only on the output-feature (N) coordinate, so read it - # once per column (global index by_n_idx + col) and reuse across the - # four M rows below. - if const_expr(has_bias): - bias_value = load_bias(by_n_idx + col) + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K64 update is two in-place K32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): + """Issue one K32 slice for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for mi in range_constexpr(4): + a_k32 = _k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _k32_frag(b_frags[nj], k32), + ) + + def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue one K32 slice for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for mi in range_constexpr(4): + a_k32 = _k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _k32_frag(b_frags[ni], k32), + ) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 - # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing: - # value = acc [+ bias] (pre-activation) - # GELU_AUX: save pre-activation to Aux, then value = gelu(value) - value = Vec(acc)[ii] + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. if const_expr(has_bias): - value = value + bias_value + bias_value = load_bias(by_n_idx + col) + + for ii in range_constexpr(4): + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) + value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + + if const_expr(has_gelu): + aux_val = value + if const_expr(output_dtype != torch.float32): + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) - if const_expr(has_gelu): - aux_val = value if const_expr(output_dtype != torch.float32): - aux_val = aux_val.to(output_fx_dtype) - aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) - fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) - value = gelu_tanh(value) - - if const_expr(output_dtype != torch.float32): - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) - - # Explicit register coordinates for HK-style four-quadrant mapping. - # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions - # inside each 128x128 quadrant: - # cA: (warp_m, warp_n) - # cB: (warp_m, warp_n + 2) - # cC: (warp_m + 2, warp_n) - # cD: (warp_m + 2, warp_n + 2) - reg_k_col0 = lane_div_16 * 16 - reg_k_col1 = 64 + lane_div_16 * 16 - - # Every fragment row differs only by multiples of 16, so row % 16 is - # always lane_mod_16. Hoist the logical->physical XOR mapping once. - _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) - _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) - - reg_subtile_m_idx0 = wave_id // 2 - reg_subtile_n_idx0 = wave_id % 2 - - reserve_pinned_accumulators() - zero_pinned_accumulators() - - def load_b_subtile_ni_regs(lds_b, sn, ni): - return _load_b_ni( - load_transposed_frag, - load_b_frag, - lds_b, - sn, - ni, - reg_subtile_n_idx0, - lane_mod_16, - ) - - def load_b_subtile_regs(lds_b, sn): - return ( - load_b_subtile_ni_regs(lds_b, sn, 0), - load_b_subtile_ni_regs(lds_b, sn, 1), - load_b_subtile_ni_regs(lds_b, sn, 2), - load_b_subtile_ni_regs(lds_b, sn, 3), - ) + value = value.to(output_fx_dtype) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + return _load_b_ni( + load_transposed_frag, + load_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) - def load_a_subtile_mi_half(lds_a, sm, mi, half): - return _load_a_half( - load_transposed_frag_half, - load_frag_half_at_byte_base, - lds_a, - sm, - mi, - half, - reg_subtile_m_idx0, - lane_mod_16, - ) + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) - def load_a_subtile_mi_regs(lds_a, sm, mi): - x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) - x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) - return pack_frag_halves(x0, x1) - - def load_a_subtile_regs(lds_a, sm): - return ( - load_a_subtile_mi_regs(lds_a, sm, 0), - load_a_subtile_mi_regs(lds_a, sm, 1), - load_a_subtile_mi_regs(lds_a, sm, 2), - load_a_subtile_mi_regs(lds_a, sm, 3), - ) + def load_a_subtile_mi_half(lds_a, sm, mi, half): + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) - def hk_one_k_with_refill( - k128, - cur_a, - cur_b, - next_a, - next_b, - refill_a, - refill_b, - a0_regs, - b0_regs, - ): + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) - # Wait only far enough for the current page; the next-page refill may remain in flight. - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): - # A-top and B-left are both carried as complete 64-row register tiles, - # so their LDS half-pages can be refilled immediately. - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + # Wait only far enough for the current page; the next-page refill may remain in flight. + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - # Q0: interleave the current tile's A-bottom LDS reads with K+2 - # refills and Q0 compute. Compute is K32-major across all 16 - # independent accumulators, eliminating the two-deep same-AGPR - # dependency chains produced by pinned_mfma(). - rocdl.sched_barrier(0) - a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a) - mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Int64((k128 + 2) * BLOCK_K) - a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) - stage_b_subtile_pass(k_refill, 0, 0, refill_b) - mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Compute is K32-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) - a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) - stage_a_subtile_pass(k_refill, 0, 1, refill_a) - # K32 slice 0 already covers K[0:32]. + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) - a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) - stage_b_subtile_pass(k_refill, 0, 1, refill_b) - # Keep this refill/LDS-read slot compute-free. + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:32]. - a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) - stage_a_subtile_pass(k_refill, 0, 2, refill_a) - mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. - a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) - stage_b_subtile_pass(k_refill, 0, 2, refill_b) - mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) - a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) - stage_a_subtile_pass(k_refill, 0, 3, refill_a) - # K32 slice 1 already covers K[32:64]. + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) - a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) - stage_b_subtile_pass(k_refill, 0, 3, refill_b) - # Keep this refill/LDS-read slot compute-free. + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[32:64]. - hot_loop_scheduler_q0_refill_a1_2n() + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. - # Retire the eight distributed A-bottom LDS reads before K+2 refills - # overwrite the current page's A-bottom half-page. Keep this wait as - # late as possible to maximize read/compute overlap. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + hot_loop_scheduler_q0_refill_a1_2n() - a10 = pack_frag_halves(a10_x0, a10_x1) - a11 = pack_frag_halves(a11_x0, a11_x1) - a12 = pack_frag_halves(a12_x0, a12_x1) - a13 = pack_frag_halves(a13_x0, a13_x1) + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 0, refill_b) - mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) - stage_a_subtile_pass(k_refill, 1, 0, refill_a) - mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) - stage_b_subtile_pass(k_refill, 1, 1, refill_b) - # K32 slice 0 already covers K[0:32]. + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) - stage_a_subtile_pass(k_refill, 1, 1, refill_a) - # Keep this refill slot compute-free. + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:32]. - stage_b_subtile_pass(k_refill, 1, 2, refill_b) - mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. - stage_a_subtile_pass(k_refill, 1, 2, refill_a) - mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) - stage_b_subtile_pass(k_refill, 1, 3, refill_b) - # K32 slice 1 already covers K[32:64]. + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) - stage_a_subtile_pass(k_refill, 1, 3, refill_a) - # Keep this refill slot compute-free. - hot_loop_scheduler_q_refill_2n() + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[32:64]. - # Leave exactly the K+2 refill and scale loads outstanding. The following - # LDS reads consume the already-ready next page, not the page being refilled. - rocdl.sched_barrier(0) - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + hot_loop_scheduler_q_refill_2n() - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - # K32 slice 0 already covers K[0:32]. + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - # K32 slice 1 already covers K[32:64]. + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - # K32 slice 0 already covers K[0:32]. + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - # K32 slice 1 already covers K[32:64]. + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) - hot_loop_scheduler_q_prefetch_4n() + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + hot_loop_scheduler_q_prefetch_4n() - return next_a0_regs, next_b0_regs + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + return next_a0_regs, next_b0_regs - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) - rocdl.sched_barrier(0) - barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + rocdl.sched_barrier(0) + barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) - hot_loop_scheduler_q_prefetch_4n() + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + hot_loop_scheduler_q_prefetch_4n() - return next_a0_regs, next_b0_regs + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): - barrier(vmcnt=0, lgkmcnt=0) + return next_a0_regs, next_b0_regs - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + barrier(vmcnt=0, lgkmcnt=0) - # Materialize the remaining final-page A/B fragments once. The - # subsequent schedule is entirely register/AGPR traffic. - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) - a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) - b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) - - # Rolling final-page epilogue. - # - # Finalize accumulators in their own physical AGPR slots, but delay - # each AGPR read/store until several independent final MFMAs have - # been issued. - # - # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, - # MFMA 4, drain 1, MFMA 5, drain 2, ... - # - # The buffer stores are only issued here; they may remain in flight - # while later MFMAs and accumulator drains continue. - FINAL_EPILOGUE_DEPTH = 4 - pending = [] - - for old_acc_idx in range_constexpr(ACCS_PER_WAVE): - subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) - a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi - b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni - - # Final MFMA remains in-place. The logical accumulator's own - # AGPR slot is unique and cannot conflict with another pending - # result, so no ad-hoc physical-slot permutation is needed. - pinned_final_mfma( - old_acc_idx, - old_acc_idx, - a_frags[a_frag_idx], - b_frags[b_frag_idx], - ) - pending.append(old_acc_idx) + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) - # Drain the oldest completed result only after enough newer - # independent MFMAs have supplied the MFMA->AGPR-read spacing. - if len(pending) == FINAL_EPILOGUE_DEPTH: - drain_acc_idx = pending.pop(0) + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: acc = read_physical_accumulator_slot(drain_acc_idx) store_acc_vector_for_logical_idx(drain_acc_idx, acc) - # Flush the final results after all final-page MFMAs have issued. - for drain_acc_idx in pending: - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) - # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) - - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - a0_regs = load_a_subtile_regs(lds_a0, 0) + a0_regs = load_a_subtile_regs(lds_a0, 0) - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - b0_regs = load_b_subtile_regs(lds_b0, 0) + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K64 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) - # Main HK loop: exactly one logical K64 per iteration. - # Even k consumes and refills LDS0; odd k does the same for LDS1. - for k128 in range_constexpr(NUM_K_TILES - 2): - if (k128 % 2) == 0: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a0, lds_b0, lds_a1, lds_b1, - lds_a0, - lds_b0, a0_regs, b0_regs, ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) else: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a1, lds_b1, lds_a0, lds_b0, - lds_a1, - lds_b1, a0_regs, b0_regs, ) - - # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch - # to prepare A-top/B-left for the final tile, but performs no K+2 refill. - if (NUM_K_TILES % 2) == 0: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a0, - lds_b0, - lds_a1, - lds_b1, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + # Dispatch. Dense runs one XCD-swizzled tile. Grouped runs a persistent + # loop over precomputed (pid_m, pid_n, expert) descriptors, selecting the + # schedule at compile time (``sched``): "worksteal" atomically claims the + # next global slot each iteration; "static" owns a fixed grid-stride slice + # (no atomic or LDS broadcast, deadlock-proof). + def _run_grouped_tile(slot): + pid_m = _desc_field(slot, fx.Int32(0)) + pid_n = _desc_field(slot, fx.Int32(1)) + expert = _desc_field(slot, fx.Int32(2)) + process_tile(pid_m, pid_n, expert) + + if const_expr(grouped): + if const_expr(sched == "static"): + stride = fx.Int32(gpu.grid_dim.x) + slot = fx.Int32(gpu.block_id("x")) + while slot < num_tiles: + _run_grouped_tile(slot) + slot = slot + stride + else: + # Single persistent loop (mega_moe_stage1 idiom): claim the next + # slot at the TOP (atomic + barrier + LDS broadcast), run the tile + # only when in range, derive the continue flag from that same + # broadcast slot -- one barrier site, workgroup-uniform condition. + active = fx.Int32(1) == fx.Int32(1) + while active: + slot = _claim_next_slot() + in_range = slot < num_tiles + if in_range: + _run_grouped_tile(slot) + active = in_range else: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a1, - lds_b1, - lds_a0, - lds_b0, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + if const_expr(use_xcd_remap): + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + process_tile(pid_m, pid_n, fx.Int32(0)) @flyc.jit def launch_gemm( @@ -1222,10 +1391,20 @@ def launch_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, stream: fx.Stream = fx.Stream(None), ): - # The integration only dispatches aligned shapes; no partial-tile masking exists. + # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU + # persistent grid that drains the descriptor queue. Placeholder work + # tensors for the dense path are allocated host-side in ``doGemm``. grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(group or wgrad): + grid_x = get_num_cus() kernel_gemm( A, B, @@ -1234,6 +1413,12 @@ def launch_gemm( Aux, c_m, c_n, + num_groups, + work_queue, + work_counter, + b_data_ptrs, + a_data_ptrs, + num_work, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1248,6 +1433,10 @@ def _cached_launch( mfma_suffix: str, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, + act_ptrs: bool = False, ): return _compile_kernel( K, @@ -1256,6 +1445,10 @@ def _cached_launch( mfma_suffix, use_xcd_remap=use_xcd_remap, epilogue=epilogue, + group=group, + sched=sched, + wgrad=wgrad, + act_ptrs=act_ptrs, ) @@ -1279,17 +1472,17 @@ def _half_prec_matmul( ``input_dtype`` is the required FP16/BF16 operand dtype and ``label`` is the human-readable kernel name used in error messages. + + Only the checks ``doGemm`` does not repeat are performed here: the + layout-specific physical operand shapes, operand/output contiguity, and + device agreement. Layout validity, operand dtype, output dtype, and tiling + are validated by ``doGemm`` against the runtime geometry, so they are not + duplicated in this shim. """ - if layout not in ("TN", "NN", "NT"): - raise ValueError(f"Unsupported {label} layout: {layout}") if a.ndim != 2 or b.ndim != 2: raise ValueError( f"FlyDSL {label} expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" ) - if a.dtype != input_dtype or b.dtype != input_dtype: - raise TypeError( - f"FlyDSL {label} GEMM expects {input_dtype} operands, got A={a.dtype}, B={b.dtype}" - ) if not a.is_contiguous() or not b.is_contiguous(): raise FlyDSLUnsupportedError( f"FlyDSL {label} {layout} requires original contiguous row-major " @@ -1306,6 +1499,8 @@ def _half_prec_matmul( "NN": ((m, k), (k, n)), "NT": ((k, m), (k, n)), } + if layout not in expected_shapes: + raise ValueError(f"Unsupported {label} layout: {layout}") expected_a, expected_b = expected_shapes[layout] if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: raise ValueError( @@ -1316,10 +1511,6 @@ def _half_prec_matmul( if tuple(c.shape) != (m, n): raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError( - f"FlyDSL {label} output must be float16, bfloat16, or float32, got {c.dtype}" - ) if a.device != b.device or a.device != c.device: raise ValueError( f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" @@ -1424,6 +1615,17 @@ def doGemm( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: torch.Tensor = None, + work_counter: torch.Tensor = None, + b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, + num_work: int = 0, + sched: str = "worksteal", + wgrad: bool = False, + a_shape=None, + b_shape=None, ): """Launch one cached K/output/layout-specialized half-precision core. @@ -1435,10 +1637,14 @@ def doGemm( NT: A backing [K,M], B backing [K,N] NN/NT orientation is implemented by compile-time global addressing and - ``ds_read_b64_tr_b16`` only. + ``ds_read_b64_tr_b16`` only. With ``a_data_ptrs`` (``b_data_ptrs`` in wgrad) the + operand is one representative expert and ``a_shape`` (``b_shape``) the packed + operand shape. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported {label} layout: {layout}") + if group and b_data_ptrs is None: + raise FlyDSLUnsupportedError(f"grouped {label} GEMM requires per-expert weight pointers") M_runtime = int(m) N_runtime = int(n) @@ -1460,10 +1666,23 @@ def doGemm( block_k=_BLOCK_K, label=f"{label} GEMM", ) - require_launch_size(f"{label} GEMM", ("A", A), ("B", B), ("C", C)) + # With per-expert pointers, A (and wgrad's B) is one representative expert, but + # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit + # offsets, so the bound is on the packed extent. + a_rows, a_cols = a_shape if a_shape is not None else A.shape + b_rows, b_cols = b_shape if b_shape is not None else B.shape + require_launch_size( + f"{label} GEMM", + ("A", A, a_rows * a_cols * A.element_size()), + ("B", B, b_rows * b_cols * B.element_size()), + ("C", C), + ) - if tuple(C.shape) != (M_runtime, N_runtime): - raise ValueError(f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}") + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is + # the unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + if tuple(C.shape) != expected_c: + raise ValueError(f"C shape {tuple(C.shape)} != expected {expected_c}") if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported {label} epilogue: {epilogue}") @@ -1497,16 +1716,8 @@ def doGemm( raise ValueError(f"{label} epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = torch.cuda.current_stream() + stream = grouped_utils.current_stream_handle(A.device) - launch = _cached_launch( - K_runtime, - C.dtype, - layout, - _MFMA_SUFFIX[input_dtype], - bool(use_xcd_remap), - epilogue, - ) # Preserve the original validated byte-addressed G2L path. These are # metadata-only dtype/flatten views of the already-contiguous row-major # tensors selected by gemm_wrappers.py; no transpose or copy is performed. @@ -1514,16 +1725,67 @@ def doGemm( B_arg = B.view(torch.uint8).view(-1) C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # The placeholders are never dereferenced by the kernel, so reuse one cached + # buffer per (device, dtype) instead of re-allocating each call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + Aux_arg = placeholder_arg(A.device, C.dtype) + + # Grouped runtime tensors are required by the kernel signature; the dense + # compiled kernel never dereferences them but still needs placeholder buffers. + if group or wgrad: + # num_work is the kernel's drain bound: it must be the queue's descriptor count. + if work_queue.numel() != 3 * int(num_work): + raise ValueError( + f"num_work={num_work} does not match work_queue of {work_queue.numel()} " + "int32 (3 per descriptor)" + ) + work_queue_arg = work_queue.contiguous().view(-1) + work_counter_arg = work_counter.contiguous().view(-1) + else: + work_queue_arg = placeholder_arg(A.device, torch.int32) + work_counter_arg = placeholder_arg(A.device, torch.int32) + + # Per-expert weight base pointers (grouped TN/NN; wgrad's pre-shifted X with + # act_ptrs): an int64[G] device array read in-kernel so the operand needs no + # host-side concat. Dense passes a placeholder that is never read. + if b_data_ptrs is not None: + b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) + else: + b_data_ptrs_arg = placeholder_arg(A.device, torch.int64) + # act_ptrs also selects how the kernel reads B: outside grouped fwd/dgrad, B is + # read through b_data_ptrs only when act_ptrs is set. So activation pointers are + # grouped-only, and wgrad takes the dY and X pointer arrays together or not at all. + if a_data_ptrs is not None and not (group or wgrad): + raise ValueError(f"{label}: per-expert activation pointers require a grouped launch") + if wgrad and (a_data_ptrs is None) != (b_data_ptrs is None): + raise ValueError(f"{label} wgrad: a_data_ptrs and b_data_ptrs must be given together") + # Per-expert activation base pointers (see act_ptrs in _compile_kernel). + act_ptrs = a_data_ptrs is not None + if act_ptrs: + a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) + else: + a_data_ptrs_arg = placeholder_arg(A.device, torch.int64) - launch( + launch_key = ( + "half_prec_gemm", + K_runtime, + C.dtype, + layout, + _MFMA_SUFFIX[input_dtype], + bool(use_xcd_remap), + epilogue, + group, + sched, + wgrad, + act_ptrs, + ) + launch_args = ( A_arg, B_arg, C_arg, @@ -1531,5 +1793,30 @@ def doGemm( Aux_arg, M_runtime, N_runtime, - stream=stream, + num_groups, + work_queue_arg, + work_counter_arg, + b_data_ptrs_arg, + a_data_ptrs_arg, + int(num_work), + stream, ) + # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are + # runtime scalars, so one compiled launcher per launch_key serves all shapes. + compiled_launch( + launch_key, + _cached_launch( + K_runtime, + C.dtype, + layout, + _MFMA_SUFFIX[input_dtype], + bool(use_xcd_remap), + epilogue, + group, + sched, + wgrad, + act_ptrs, + ), + launch_args, + )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 375f11d5a..f2ac5b3d1 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -34,8 +34,15 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + get_num_cus, + placeholder_arg, + require_block_tiling, + require_launch_size, +) from .fp8_gemm_utils import ( G2SLoader, S2RLoader, @@ -78,12 +85,19 @@ def _compile_mx32_scale_pack_kernel( source_colwise: bool, stride0: int, stride1: int, + per_expert: bool = False, ): """Build one fused scale packing kernel. One GPU thread produces one final ``uint32`` word in the GEMM-consumed ``[K/128, dim]`` layout. There is no intermediate ``scale_iter`` tensor and no eager PyTorch shift/index/OR kernels. + + ``per_expert`` packs the concatenation of per-expert scale tensors without + building it: the source is an int64[G] array of per-expert base pointers, each + shifted to its expert's offset in the concatenated layout (so the offsets below + are unchanged), plus an int32 table mapping each 64-row tile (rowwise) or K128 + block (columnwise) to its expert. """ if dim % 64 != 0: raise FlyDSLUnsupportedError(f"Scale outer dimension={dim} must be a multiple of 64") @@ -98,69 +112,108 @@ def _compile_mx32_scale_pack_kernel( ) # Select source addressing before FlyDSL captures the kernel. The emitted - # rowwise and columnwise binaries contain no runtime orientation branch. + # rowwise and columnwise compiled kernels contain no runtime orientation branch. if source_colwise: def _source_offset(source_k32, source_row): # Logical source is [K/32, dim], but the underlying TE tensor may # be a non-contiguous view. Strides are in uint8 elements. - return source_k32 * fx.Index(stride0) + source_row * fx.Index(stride1) + return source_k32 * fx.Int64(stride0) + source_row * fx.Int64(stride1) else: def _source_offset(source_k32, source_row): # Logical source is [dim, K/32], with arbitrary positive strides. - return source_row * fx.Index(stride0) + source_k32 * fx.Index(stride1) + return source_row * fx.Int64(stride0) + source_k32 * fx.Int64(stride1) - @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) - def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): - # Address both buffers as flat 1-D arrays so a linear slice coordinate - # maps to base+off, matching the legacy buffer_load semantics. Building - # buffer tensors directly from the 2-D src/dst would make logical_divide - # walk the layout column-major and corrupt every strided access. - src_flat = fx.rocdl.make_buffer_tensor( - fx.Tensor(fx.make_view(fx.get_iter(src), fx.make_layout(dim * qk, 1))), - max_size=True, + def _word_coords(): + # This thread's output word and the source coordinates of its four bytes. + linear = fx.Int64(fx.block_idx.x) * fx.Int64(_SCALE_PACK_THREADS) + fx.Int64( + gpu.thread_id("x") ) + k128 = linear // fx.Int64(dim) + dst_row = linear % fx.Int64(dim) + + row_within_16 = dst_row % fx.Int64(16) + k_subgroup = (dst_row // fx.Int64(16)) % fx.Int64(4) + tile = dst_row // fx.Int64(64) + source_k32 = k128 * fx.Int64(4) + k_subgroup + return linear, k128, tile, row_within_16, source_k32 + + def _load_scale_byte(src_div, coords, group): + _, _, tile, row_within_16, source_k32 = coords + source_row = tile * fx.Int64(64) + fx.Int64(group * 16) + row_within_16 + off = _source_offset(source_k32, source_row) + scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) + fx.copy(scale_load_atom, fx.slice(src_div, (None, fx.Int32(off))), reg) + value_u8 = fx.memref_load_vec(reg)[0] + # The byte load widens via sext, so mask to the low 8 bits to + # preserve the raw E8M0 byte for scales >= 0x80 before packing. + return fx.Int32(value_u8) & fx.Int32(0xFF) + + def _store_word(dst, src_div, coords): dst_flat = fx.rocdl.make_buffer_tensor( fx.Tensor(fx.make_view(fx.get_iter(dst), fx.make_layout(total_words, 1))), max_size=True, ) - src_div = fx.logical_divide(src_flat, fx.make_layout(1, 1)) dst_div = fx.logical_divide(dst_flat, fx.make_layout(1, 1)) - scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) scale_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) - - linear = fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) + fx.Index( - gpu.thread_id("x") - ) - k128 = linear // fx.Index(dim) - dst_row = linear % fx.Index(dim) - - row_within_16 = dst_row % fx.Index(16) - k_subgroup = (dst_row // fx.Index(16)) % fx.Index(4) - tile = dst_row // fx.Index(64) - source_k32 = k128 * fx.Index(4) + k_subgroup - - def load_scale_byte(group): - source_row = tile * fx.Index(64) + fx.Index(group * 16) + row_within_16 - off = _source_offset(source_k32, source_row) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) - fx.copy(scale_load_atom, fx.slice(src_div, (None, fx.Int32(off))), reg) - value_u8 = fx.memref_load_vec(reg)[0] - # The byte load widens via sext, so mask to the low 8 bits to - # preserve the raw E8M0 byte for scales >= 0x80 before packing. - return fx.Int32(value_u8) & fx.Int32(0xFF) - - b0 = load_scale_byte(0) - b1 = load_scale_byte(1) - b2 = load_scale_byte(2) - b3 = load_scale_byte(3) + b0 = _load_scale_byte(src_div, coords, 0) + b1 = _load_scale_byte(src_div, coords, 1) + b2 = _load_scale_byte(src_div, coords, 2) + b3 = _load_scale_byte(src_div, coords, 3) packed = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) reg_i32 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) fx.memref_store_vec(Vec.filled(1, packed, fx.Int32), reg_i32) + linear = coords[0] fx.copy(scale_store_atom, reg_i32, fx.slice(dst_div, (None, fx.Int32(linear)))) + if per_expert: + + @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) + def kernel_pack_mx32_scales(src_ptrs: fx.Tensor, expert_table: fx.Tensor, dst: fx.Tensor): + coords = _word_coords() + _, k128, tile, _, _ = coords + # A wave's 64 consecutive words share one 64-row tile (dim % 64 == 0) and one + # K128 block, so one expert: make its base wave-uniform and rebase a buffer + # view, keeping the buffer loads of the packed path. + block = k128 if source_colwise else tile + expert = fx.Int32(grouped_utils.load_i32_at(expert_table, block)) + base = grouped_utils.make_uniform_i64(grouped_utils.read_ptr_at(src_ptrs, expert)) + src_div = grouped_utils.rebase_buffer_view( + base, fx.Uint8.ir_type, fx.make_layout(dim * qk, 1) + ) + _store_word(dst, src_div, coords) + + @flyc.jit + def launch_pack_mx32_scales( + src_ptrs: fx.Tensor, + expert_table: fx.Tensor, + dst: fx.Tensor, + stream: fx.Stream = fx.Stream(None), + ): + kernel_pack_mx32_scales(src_ptrs, expert_table, dst).launch( + grid=(total_words // _SCALE_PACK_THREADS, 1, 1), + block=(_SCALE_PACK_THREADS, 1, 1), + stream=stream, + ) + + return launch_pack_mx32_scales + + @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) + def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): + # Address both buffers as flat 1-D arrays so a linear slice coordinate + # maps to base+off, matching the legacy buffer_load semantics. Building + # buffer tensors directly from the 2-D src/dst would make logical_divide + # walk the layout column-major and corrupt every strided access. + src_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(fx.get_iter(src), fx.make_layout(dim * qk, 1))), + max_size=True, + ) + src_div = fx.logical_divide(src_flat, fx.make_layout(1, 1)) + _store_word(dst, src_div, _word_coords()) + @flyc.jit def launch_pack_mx32_scales( src: fx.Tensor, @@ -183,9 +236,10 @@ def _cached_mx32_scale_pack_launch( source_colwise: bool, stride0: int, stride1: int, + per_expert: bool = False, ): - """Cache orientation-and-stride-specialized fused pack binaries.""" - return _compile_mx32_scale_pack_kernel(dim, qk, source_colwise, stride0, stride1) + """Cache orientation-and-stride-specialized fused pack compiled kernels.""" + return _compile_mx32_scale_pack_kernel(dim, qk, source_colwise, stride0, stride1, per_expert) def pack_mx32_scales_for_hk( @@ -228,23 +282,95 @@ def pack_mx32_scales_for_hk( device=scales_u8.device, ) if stream is None: - stream = torch.cuda.current_stream(scales_u8.device) + stream = grouped_utils.current_stream_handle(scales_u8.device) stride0, stride1 = (int(x) for x in scales_u8.stride()) - _cached_mx32_scale_pack_launch( - dim, - qk, - bool(source_colwise), - stride0, - stride1, - )( - scales_u8, - packed, - stream=stream, + pack_key = ("mx32_scale_pack", dim, qk, bool(source_colwise), stride0, stride1) + pack_args = (scales_u8, packed, stream) + compiled_launch( + pack_key, + _cached_mx32_scale_pack_launch(dim, qk, bool(source_colwise), stride0, stride1), + pack_args, + )(*pack_args) + return packed + + +def pack_mx32_scales_for_hk_experts( + scales, + *, + source_colwise: bool = False, + stream=None, +) -> torch.Tensor: + """:func:`pack_mx32_scales_for_hk` of ``torch.cat(scales)``, without the concat. + + ``scales`` are per-expert uint8 E8M0 tensors stacked on dim 0: rowwise + ``[rows_e, K/32]`` or columnwise ``[K32_e, dim]``. The packing kernel reads each + expert in place through ``grouped_gemm_utils.stacked_operand`` pointers. Falls + back to packing a real concat when the tensors cannot be addressed in place. + """ + # stacked_operand requires one common trailing shape, so rank 2 holds for all. + stacked = ( + grouped_utils.stacked_operand(scales, scales[0].device) if scales[0].dim() == 2 else None ) + if stacked is None: + return pack_mx32_scales_for_hk( + torch.cat(scales, dim=0), source_colwise=source_colwise, stream=stream + ) + representative, src_ptrs = stacked + if representative.dtype != torch.uint8: + raise TypeError(f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {representative.dtype}") + + extents = [t.shape[0] for t in scales] + total = sum(extents) + inner = representative.shape[1] + if source_colwise: + qk, dim = total, inner + # Experts stack on K/32; each thread's bytes share one K128 block (4 rows). + block = 4 + else: + dim, qk = total, inner + # Experts stack on dim; each thread's bytes share one 64-row tile. + block = 64 + expert_table = grouped_utils.block_expert_table(extents, block, representative.device) + + packed = torch.empty((qk // 4, dim), dtype=torch.int32, device=representative.device) + if stream is None: + stream = grouped_utils.current_stream_handle(representative.device) + # The shifted pointers address the contiguous concatenation, so its strides apply. + stride0, stride1 = inner, 1 + pack_key = ("mx32_scale_pack_experts", dim, qk, bool(source_colwise), stride0, stride1) + pack_args = (src_ptrs, expert_table, packed, stream) + compiled_launch( + pack_key, + _cached_mx32_scale_pack_launch( + dim, qk, bool(source_colwise), stride0, stride1, per_expert=True + ), + pack_args, + )(*pack_args) return packed +def _scale_shape(scale): + # A scale operand is one tensor or a list of per-expert tensors stacked on dim 0. + if isinstance(scale, (list, tuple)): + return (sum(t.size(0) for t in scale), scale[0].size(1)) + return tuple(scale.shape) + + +def _pack_scales(scale, *, source_colwise, stream): + if isinstance(scale, (list, tuple)): + return pack_mx32_scales_for_hk_experts(scale, source_colwise=source_colwise, stream=stream) + return pack_mx32_scales_for_hk(scale, source_colwise=source_colwise, stream=stream) + + +def _scale_dtype(scale): + return scale[0].dtype if isinstance(scale, (list, tuple)) else scale.dtype + + +def _scale_device(scale): + return scale[0].device if isinstance(scale, (list, tuple)) else scale.device + + def _compile_kernel( K: int, a_fp8_dtype: torch.dtype, @@ -252,6 +378,10 @@ def _compile_kernel( output_dtype: torch.dtype, layout: str, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, + act_ptrs: bool = False, ): """Build one compile-time-specialized TN, NN, or NT kernel. @@ -359,7 +489,7 @@ def _a_leading_dim(c_m): return c_m def _a_global_base(k_base, subtile, c_m, bx_m_idx): - return k_base * fx.Index(c_m) + bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + return k_base * fx.Int64(c_m) + bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2)) def _load_a_half( load_transposed_frag_half, @@ -371,11 +501,11 @@ def _load_a_half( reg_subtile_m_idx0, lane_mod_16, ): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) local_m_tile = ( - subtile_m_idx * fx.Index(SUBTILE_M) - + fx.Index(mi * MFMA_M) - - fx.Index(sm * (BLOCK_M // 2)) + subtile_m_idx * fx.Int64(SUBTILE_M) + + fx.Int64(mi * MFMA_M) + - fx.Int64(sm * (BLOCK_M // 2)) ) return load_transposed_frag_half( lds_a[sm], @@ -391,7 +521,7 @@ def _a_leading_dim(c_m): def _a_global_base(k_base, subtile, c_m, bx_m_idx): del c_m - return (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base + return (bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2))) * fx.Int64(K) + k_base def _load_a_half( load_transposed_frag_half, @@ -404,10 +534,10 @@ def _load_a_half( lane_mod_16, ): del load_transposed_frag_half - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Int64(BLOCK_K) return load_frag_half_at_byte_base( lds_a[sm], row_byte_base, @@ -420,7 +550,7 @@ def _b_leading_dim(c_n): return c_n def _b_global_base(k_base, subtile, c_n, by_n_idx): - return k_base * fx.Index(c_n) + by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) + return k_base * fx.Int64(c_n) + by_n_idx + fx.Int64(subtile * (BLOCK_N // 2)) def _load_b_ni( load_transposed_frag, @@ -432,11 +562,11 @@ def _load_b_ni( lane_mod_16, ): del load_normal_b_frag, lane_mod_16 - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) + subtile_n_idx * fx.Int64(SUBTILE_N) + + fx.Int64(ni * MFMA_N) + - fx.Int64(sn * (BLOCK_N // 2)) ) return load_transposed_frag(lds_b[sn], local_n_tile) @@ -448,7 +578,7 @@ def _b_leading_dim(c_n): def _b_global_base(k_base, subtile, c_n, by_n_idx): del c_n - return (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base + return (by_n_idx + fx.Int64(subtile * (BLOCK_N // 2))) * fx.Int64(K) + k_base def _load_b_ni( load_transposed_frag, @@ -460,8 +590,8 @@ def _load_b_ni( lane_mod_16, ): del load_transposed_frag - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 return load_normal_b_frag(lds_b, b_row_addr, sn) if (not a_transpose_read) or (not b_transpose_read): @@ -491,6 +621,9 @@ class SharedStorage: b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + # Persistent grouped path: one-int32 scratch to broadcast the atomically + # claimed queue slot from thread 0 to the whole workgroup. + claim_slot: fx.Array[fx.Int32, 1, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -503,7 +636,17 @@ def kernel_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, ): + # Grouped drain bound: the exact number of descriptors the host wrote to + # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. + num_tiles = num_work + lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) lds_a1 = (lds.a1_0, lds.a1_1) @@ -522,972 +665,1204 @@ def kernel_gemm( bs_div = fx.logical_divide(bs_rsrc, fx.make_layout(1, 1)) scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) tx = gpu.thread_id("x") + tx_i32 = fx.Int32(tx) num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) - - bx_m = pid_m * BLOCK_M - by_n = pid_n * BLOCK_N - - # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # Compile-time global leading dimensions: - # normal source [X,K] -> leading dimension K - # transpose source [K,X] -> leading dimension X - a_leading_dim = _a_leading_dim(c_m) - b_leading_dim = _b_leading_dim(c_n) - - # gl_off_a/gl_off_b: per-lane lists of LOAD_PASSES_HALF swizzled flat - # global offsets, indexed by DMA pass. gl_off_a[step] is this lane's - # static 16-byte source for pass `step`; G2SLoader adds the dynamic - # K-tile base as soffset. Offsets are pre-swizzled so bytes land in the - # bank-conflict-free LDS slots the MFMA read (S2RLoader) expects. - gl_off_a = compute_global_swizzle( - lane, - wave_id, - a_leading_dim, - LOAD_PASSES_HALF, - preshuffled=False, - ) - gl_off_b = compute_global_swizzle( - lane, - wave_id, - b_leading_dim, - LOAD_PASSES_HALF, - preshuffled=False, - ) - a_g2s = G2SLoader( - a_div, - gl_off_a, - LOAD_PASSES_HALF, - a_f8_ir_t, - wave_id, - ) - b_g2s = G2SLoader( - b_div, - gl_off_b, - LOAD_PASSES_HALF, - b_f8_ir_t, - wave_id, - ) - s2r = S2RLoader(fx.Int32(0), 1) - - layout_lane16 = fx.make_layout((4, 16), (16, 1)) - coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) - lane_div_16 = fx.get(coord_lane16, 0) - lane_mod_16 = fx.get(coord_lane16, 1) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - # Bias is a length-N fp32 vector indexed by the output-feature (N) - # coordinate and broadcast across the M/token rows. Only staged when - # the epilogue requests it; DEFAULT receives a dummy 1-element tensor. - # const_expr folds this compile-time flag at trace time so the setup - # (and load_bias) are inlined into the kernel scope with no runtime - # dispatch branch. - if const_expr(has_bias): - gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) - bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) - bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - - def load_bias(col): - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) - return fx.memref_load_vec(reg)[0] - - # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N - # output so the backward pass can recompute the GELU gradient. Same tile - # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. - if const_expr(has_gelu): - gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) - aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) - aux_store_atom = fx.make_copy_atom( + # Both the grouped fwd/dgrad path and the wgrad specialization use the same + # persistent descriptor-queue driver; ``grouped`` gates that shared setup. + grouped = group or wgrad + + # Persistent work-stealing driver. + # + # Dense path (group is False): the XCD-swizzled block id owns exactly one + # (pid_m, pid_n) output tile; the body below runs once. + # + # Grouped path (group is True, packed TN, const_expr-folded so the dense + # compiled kernel is untouched): the launch is a fixed num-CU grid of persistent + # workgroups. Each workgroup claims a queue slot by atomically fetch-adding + # the shared counter (``work_counter[0]``), reads the precomputed + # ``(pid_m, pid_n, expert)`` descriptor at that slot from ``work_queue``, and + # keeps claiming until the queue drains -- so faster CUs process more tiles. + # The host bakes the tile->expert mapping into the descriptor list (no + # in-kernel search), and may permute it (XCD-aware, group-balanced, ...) + # without any kernel change. B is packed [G*N, K] and B-scale [K128, G*N], + # so B reads use the expert-shifted column ``expert*N + n`` with stride + # ``G*N`` while A / A-scale / C / bias keep absolute-row / local-N coords. + if const_expr(grouped): + wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) + wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + + # Grouped fwd/dgrad: the weights arrive as G independently-allocated + # per-expert slabs (not one packed buffer), so B is addressed from a + # per-expert base pointer read from ``b_data_ptrs`` rather than an + # expert-shifted offset into a packed B. Reuse ``gB``'s static layout + # (compile-time cosize) and swap only the base per tile. wgrad reads B by + # pointer only with act_ptrs; dense never enters here. + b_ptr_layout = fx.get_layout(gB) + # act_ptrs: the per-expert activation operand (A; and in wgrad also B, via + # b_data_ptrs) stays in its own per-expert allocation instead of a host-side + # concat. The host passes each expert's base pointer pre-shifted by that + # expert's packed offset, so the packed addressing below is unchanged; only + # the buffer base moves per tile. + if const_expr(act_ptrs): + a_ptr_layout = fx.get_layout(gA) + + # LDS broadcast scratch (one int32), following the persistent-kernel + # idiom in mega_moe_stage1: recast the LDS slot to an int32 iterator and + # a 1-element view. Thread 0 writes via ptr_store; every thread reads via + # the view after a workgroup barrier. (Descriptor reads reuse the int32 + # ``scale_ld_atom`` as their copy atom.) + claim_iter, claim_view = grouped_utils.make_claim_broadcast(lds.claim_slot) + + def _claim_next_slot(): + # One atomic per workgroup: thread 0 fetch-adds the shared counter + # and stores the pre-update value into the LDS scratch; a barrier + # then broadcasts it so every thread returns the same slot. The + # ``if tx==0`` control flow must stay inline (AST-rewritten here). + if tx_i32 == fx.Int32(0): + old = llvm.AtomicRMWOp( + llvm.AtomicBinOp.add, + grouped_utils.global_i32_ptr(work_counter, fx.Int32(0)), + fx.as_ir_value(fx.Int32(1)), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + fx.ptr_store(fx.Vector.from_elements([fx.Int32(old)], fx.Int32), claim_iter) + fx.barrier() + return fx.Vector(claim_view.load())[0] + + def _desc_field(slot, field): + return grouped_utils.read_desc_field(wq_div, scale_ld_atom, slot, field) + + def process_tile(pid_m, pid_n, expert): + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) + + # Per-expert offsets. Grouped fwd/dgrad read B data per expert through + # b_data_ptrs (see b_div_tile below), so B-data addressing is local and + # b_expert_k_off is 0. Only the B scales stay stacked over the G experts, + # on B's first logical axis, which is a different kernel axis per layout: + # + # TN (forward, C = A @ B^T): B-scale stacks experts on the OUTPUT-N + # axis, so the scale column is by_n_b_idx = expert*N + n and the + # scale row stride grows to G*N (c_n_bs). + # + # NN (dgrad, C = A @ B where A=dY[M,N_out], B=W[N_out,K_in]): the + # kernel contraction axis is N_out, and the B-scale stacks experts + # on it, so the scale k128 term shifts by expert*(N_out/128) + # (bs_expert_k128_off). The scale column and stride stay per-expert. + # + # Dense (group is False) zeroes every expert term. + # + # wgrad (NT, dW = dY^T @ X): a separate specialization. Both operands + # transpose-read and BOTH contract over the token axis M_tok (the + # kernel's compile-time K = the per-expert padded token count M_pad). + # Experts are padded to a uniform M_pad and packed on the contraction + # axis, so both a (dY [G*M_pad, N_out]) and b (X [G*M_pad, K_in]) + # shift by expert*K along k_base (a_expert_k_off / b_expert_k_off) and + # expert*NUM_K_TILES along their scale k128 terms. The output is packed + # [G*N_out, K_in] on the OUTPUT-M axis, so the store row gets an extra + # expert*N_out (= expert*c_m) offset (store_m_expert_off); the a/b + # operand row/col coords stay per-expert (local pid_m over N_out). + if const_expr(wgrad): + by_n_b_idx = by_n_idx + c_n_bs = c_n + a_expert_k_off = fx.Int64(expert * fx.Int32(K)) + b_expert_k_off = fx.Int64(expert * fx.Int32(K)) + as_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + store_m_expert_off = fx.Int64(expert * fx.Int32(c_m)) + elif const_expr(group): + if const_expr(b_transpose_read): + # NN: the B-scale [G*N_out/128, K_in] stacks experts on the + # contraction axis (compile-time K = N_out), so it shifts by + # expert*NUM_K_TILES. B data is read per expert (no shift). + by_n_b_idx = by_n_idx + c_n_bs = c_n + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) + as_expert_k128_off = fx.Int32(0) + bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + store_m_expert_off = fx.Int64(0) + else: + # TN: weights are per-expert slabs addressed via a base-pointer + # swap (b_div_tile below), so B's column coord stays LOCAL + # (by_n_idx) -- the expert shift lives in the base, not the + # column. The B-scale still packs [K128, G*N], so its row stride + # grows to G*N (c_n_bs) and column folds in expert (handled at + # by_n_b_idx for the scale read; B-data uses by_n_idx directly + # via the per-expert base). + by_n_b_idx = fx.Int64(expert * fx.Int32(c_n) + by_n) + c_n_bs = num_groups * c_n + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) + as_expert_k128_off = fx.Int32(0) + bs_expert_k128_off = fx.Int32(0) + store_m_expert_off = fx.Int64(0) + else: + by_n_b_idx = by_n_idx + c_n_bs = c_n + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) + as_expert_k128_off = fx.Int32(0) + bs_expert_k128_off = fx.Int32(0) + store_m_expert_off = fx.Int64(0) + + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # Compile-time global leading dimensions: + # normal source [X,K] -> leading dimension K + # transpose source [K,X] -> leading dimension X + a_leading_dim = _a_leading_dim(c_m) + b_leading_dim = _b_leading_dim(c_n) + + # gl_off_a/gl_off_b: per-lane lists of LOAD_PASSES_HALF swizzled flat + # global offsets, indexed by DMA pass. gl_off_a[step] is this lane's + # static 16-byte source for pass `step`; G2SLoader adds the dynamic + # K-tile base as soffset. Offsets are pre-swizzled so bytes land in the + # bank-conflict-free LDS slots the MFMA read (S2RLoader) expects. + gl_off_a = compute_global_swizzle( + lane, + wave_id, + a_leading_dim, + LOAD_PASSES_HALF, + preshuffled=False, + ) + gl_off_b = compute_global_swizzle( + lane, + wave_id, + b_leading_dim, + LOAD_PASSES_HALF, + preshuffled=False, + ) + # Issue both per-expert base-pointer loads before either is made uniform, so a + # tile waits on one memory round trip rather than two. + if const_expr(act_ptrs): + a_base_raw = grouped_utils.read_ptr_at(a_data_ptrs, expert) + if const_expr(group or act_ptrs): + b_base_raw = grouped_utils.read_ptr_at(b_data_ptrs, expert) + if const_expr(act_ptrs): + a_base_i64 = grouped_utils.make_uniform_i64(a_base_raw) + a_div_tile = grouped_utils.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) + else: + a_div_tile = a_div + a_g2s = G2SLoader( + a_div_tile, + gl_off_a, + LOAD_PASSES_HALF, + a_f8_ir_t, + wave_id, + ) + # Grouped B-data source: rebase onto this tile's expert weight (the + # per-expert base pointer already selects the expert), so in-tile + # addressing is local (column by_n_idx, no contraction shift) for both + # TN and NN. wgrad bases are pre-shifted (act_ptrs), so they keep the + # packed column and contraction shift. Dense keeps the shared ``b_div``. + if const_expr(group): + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_data_col = by_n_idx + b_k_off = fx.Int64(0) + elif const_expr(act_ptrs): + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off + else: + b_div_tile = b_div + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off + b_g2s = G2SLoader( + b_div_tile, + gl_off_b, + LOAD_PASSES_HALF, + b_f8_ir_t, + wave_id, + ) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + # store_m_expert_off (grouped wgrad only) shifts the output row into this + # tile's expert block of the packed [G*N_out, K_in] weight-gradient; the + # a/b operands used local (per-expert) pid_m, so the store re-packs here. + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Int64(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), output_fx_dtype, ) - def gelu_tanh(x): - # tanh-approx GELU (matches PyTorch approximate='tanh' and the - # FlyDSL preshuffle reference), expressed through a non-positive - # exponent so exp() cannot overflow: - # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = x * x * x - y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * x * (numerator * (one_f32 / denom)) - - PIN_ACC_BASE = 0 - - def _reg_list(prefix, start, end): - return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) - - def reserve_pinned_accumulators(): - # Reserve a fixed physical AGPR bank for all accumulators. In the - # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, - # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator - # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the - # scaled MFMA accumulation in place and avoids those transfers and spills. - # - # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, - # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. - clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) - llvm.InlineAsmOp( - None, - [], - "", - clobbers, - has_side_effects=True, - ) + # Bias is a length-N fp32 vector indexed by the output-feature (N) + # coordinate and broadcast across the M/token rows. Only staged when + # the epilogue requests it; DEFAULT receives a dummy 1-element tensor. + # const_expr folds this compile-time flag at trace time so the setup + # (and load_bias) are inlined into the kernel scope with no runtime + # dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy32b() + if output_element_bytes == 4 + else fx.rocdl.BufferCopy16b() + ), + output_fx_dtype, + ) - def zero_pinned_accumulators(): - for ai in range_constexpr(ACCS_PER_WAVE * 4): + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) llvm.InlineAsmOp( None, [], - f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", - f"~{{a{PIN_ACC_BASE + ai}}}", + "", + clobbers, has_side_effects=True, ) - def _inline_asm_i32(asm_string, constraints, operands=None): - op = llvm.InlineAsmOp( - T.i32, - operands or [], - asm_string, - constraints, - has_side_effects=True, - ) - return _one_i32_result(op) - - def _one_i32_result(op): - # Accept the result attribute names exposed by the supported MLIR Python bindings. - return getattr(op, "result", getattr(op, "res", op.results[0])) - - def _to_raw_inline_asm_operand(value): - # TODO: Replace arith._to_raw once FlyDSL exposes a supported public - # API for passing wrapped values to llvm.InlineAsmOp. _to_raw is - # deprecated, but remains heavily used internally by FlyDSL. - return arith._to_raw(value) - - def read_physical_accumulator_slot(slot_idx): - acc_pin = PIN_ACC_BASE + slot_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - # As/Bs are MFMA-ready packed scale words: [K128, row] uint32. - # Each loaded dword already contains the four 16-row/16-col MFMA scale - # bytes for this lane's 64-row A/B half. The MFMA instruction selects - # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop - # byte extraction and no 0x01010101 broadcast here. - c_m_idx = fx.Index(c_m) - c_n_idx = fx.Index(c_n) - - def hot_loop_scheduler_q_refill_2n(): - # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS - # refill pass followed by two MFMAs. - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_mfma(2) - - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q0_refill_a1_2n(): - # TN/NN: one normal A-bottom LDS read per chunk. - # NT: one transpose-read A half plus the matching transpose-read - # scheduling pressure retained from the passing NT specialization. - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_dsrd(Q0_SCHED_DSRD) - rocdl.sched_mfma(2) - - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q_prefetch_4n(): - # TN/NN retain two scheduled DS reads per chunk. NT retains four - # because both carried operands use two DS_READ_TR instructions. - for _ in range_constexpr(8): - rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) - rocdl.sched_mfma(4) + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def _to_raw_inline_asm_operand(value): + # TODO: Replace arith._to_raw once FlyDSL exposes a supported public + # API for passing wrapped values to llvm.InlineAsmOp. _to_raw is + # deprecated, but remains heavily used internally by FlyDSL. + return arith._to_raw(value) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + # As/Bs are MFMA-ready packed scale words: [K128, row] uint32. + # Each loaded dword already contains the four 16-row/16-col MFMA scale + # bytes for this lane's 64-row A/B half. The MFMA instruction selects + # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop + # byte extraction and no 0x01010101 broadcast here. + c_m_idx = fx.Int64(c_m) + c_n_idx = fx.Int64(c_n) + # B-scale row stride: c_n dense, packed G*N grouped (const_expr-folded). + c_n_bs_idx = fx.Int64(c_n_bs) + + def hot_loop_scheduler_q_refill_2n(): + # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS + # refill pass followed by two MFMAs. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # TN/NN: one normal A-bottom LDS read per chunk. + # NT: one transpose-read A half plus the matching transpose-read + # scheduling pressure retained from the passing NT specialization. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(Q0_SCHED_DSRD) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # TN/NN retain two scheduled DS reads per chunk. NT retains four + # because both carried operands use two DS_READ_TR instructions. + for _ in range_constexpr(8): + rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) + rocdl.sched_mfma(4) + + rocdl.sched_barrier(0) + + def load_a_scale_row(k128, row): + # as_expert_k128_off shifts the contraction-128 base to this tile's + # expert A-scale block for grouped wgrad (both operands packed on the + # token/contraction axis); 0 for every other path. + off = (k128 + fx.Int64(as_expert_k128_off)) * c_m_idx + bx_m_idx + row + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) + return fx.memref_load_vec(reg)[0] - rocdl.sched_barrier(0) + def load_b_scale_row(k128, row): + # bs_expert_k128_off shifts the contraction-128 base to this tile's + # expert scale block for grouped NN/NT; 0 for dense and grouped TN. + off = (k128 + fx.Int64(bs_expert_k128_off)) * c_n_bs_idx + by_n_b_idx + row + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) + return fx.memref_load_vec(reg)[0] - def load_a_scale_row(k128, row): - off = k128 * c_m_idx + bx_m_idx + row - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) - fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) - return fx.memref_load_vec(reg)[0] - - def load_b_scale_row(k128, row): - off = k128 * c_n_idx + by_n_idx + row - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) - fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) - return fx.memref_load_vec(reg)[0] - - def load_a_scale_subtile(k128, sm): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(lane) - a_scale = load_a_scale_row(k128, a_row) - return (a_scale, a_scale, a_scale, a_scale) - - def load_b_scale_subtile(k128, sn): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(lane) - b_scale = load_b_scale_row(k128, b_row) - return (b_scale, b_scale, b_scale, b_scale) - - def load_scale_tile(k128): - # Load all scale VGPRs needed by this wave for this K128 tile once. - # Return order: A-top, A-bottom, B-left, B-right. - return ( - load_a_scale_subtile(k128, 0), - load_a_scale_subtile(k128, 1), - load_b_scale_subtile(k128, 0), - load_b_scale_subtile(k128, 1), - ) + def load_a_scale_subtile(k128, sm): + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(lane) + a_scale = load_a_scale_row(k128, a_row) + return (a_scale, a_scale, a_scale, a_scale) + + def load_b_scale_subtile(k128, sn): + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(lane) + b_scale = load_b_scale_row(k128, b_row) + return (b_scale, b_scale, b_scale, b_scale) + + def load_scale_tile(k128): + # Load all scale VGPRs needed by this wave for this K128 tile once. + # Return order: A-top, A-bottom, B-left, B-right. + return ( + load_a_scale_subtile(k128, 0), + load_a_scale_subtile(k128, 1), + load_b_scale_subtile(k128, 0), + load_b_scale_subtile(k128, 1), + ) - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - a_g2s.load_one( - lds_a[subtile], - fx.Int32(_a_global_base(k_base, subtile, c_m, bx_m_idx)), - pass_in_subtile, - ) + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # a_expert_k_off shifts the contraction base to this tile's expert + # A block for grouped wgrad (dY packed on the token axis); 0 for + # every other path. + a_g2s.load_one( + lds_a[subtile], + fx.Int32(_a_global_base(k_base + a_expert_k_off, subtile, c_m, bx_m_idx)), + pass_in_subtile, + ) - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - b_g2s.load_one( - lds_b[subtile], - fx.Int32(_b_global_base(k_base, subtile, c_n, by_n_idx)), - pass_in_subtile, - ) + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # Grouped fwd/dgrad select the expert via the rebased per-expert B view + # (b_g2s) with LOCAL coords: TN a local column (b_data_col), NN a zero + # contraction shift (b_k_off). Dense/wgrad use the shared view with the + # packed column by_n_b_idx and the expert contraction shift b_k_off. + b_g2s.load_one( + lds_b[subtile], + fx.Int32(_b_global_base(k_base + b_k_off, subtile, c_n, b_data_col)), + pass_in_subtile, + ) - def stage_a_subtile(k_base, subtile, lds_a): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - - def stage_b_subtile(k_base, subtile, lds_b): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) - - def load_frag_half_at_byte_base(lds_page, row_byte_base, half): - # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. - # Keeping the halves separate allows steady-state Q0 to schedule one - # A-bottom ds_read_b128 in each refill/MFMA chunk. - k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 - return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) - - def pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_frag_at_byte_base(lds_page, row_byte_base): - # Default complete-fragment path used outside the dedicated Q0 schedule. - x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) - x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) - return pack_frag_halves(x0, x1) - - def load_normal_b_frag(lds_b, local_row, half): - # Physical [N,K] page, ordinary TN-style fixed-row read. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base( - lds_b[half], - half_row * fx.Index(BLOCK_K), - ) + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_normal_b_frag(lds_b, local_row, half): + # Physical [N,K] page, ordinary TN-style fixed-row read. + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Int64(BLOCK_K), + ) - def load_transposed_frag_half(lds_page, local_x_tile, half): - # Exact inverse mapping validated by the MXFP8 NN fragment probe. - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - source_k = lane_div16_i32 * fx.Int32(16) + lane_in16_i32 // fx.Int32(2) - source_x = fx.Int32(local_x_tile) + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) - - physical_k, physical_x = swizzle_128(source_k, source_x) - base = physical_k * fx.Int32(128) + physical_x - other = base ^ fx.Int32(0x440) - immediate_offset = 0 if half == 0 else 0x2000 - - return s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=immediate_offset, - ) + def load_transposed_frag_half(lds_page, local_x_tile, half): + # Exact inverse mapping validated by the MXFP8 NN fragment probe. + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = lane_div16_i32 * fx.Int32(16) + lane_in16_i32 // fx.Int32(2) + source_x = fx.Int32(local_x_tile) + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + + physical_k, physical_x = swizzle_128(source_k, source_x) + base = physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x440) + immediate_offset = 0 if half == 0 else 0x2000 + + return s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) - def load_transposed_frag(lds_page, local_x_tile): - x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) - x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) - return pack_frag_halves(x0, x1) + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) + return pack_frag_halves(x0, x1) - def _acc_idx(subtile_id, mi, ni): - return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni - ) + def _acc_idx(subtile_id, mi, ni): + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + mi * MFMA_N_PER_SUBTILE + + ni + ) - def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): - # Fixed physical accumulator bank, visible SSA A/B/scale operands. - # acc_idx maps directly to a[PIN_ACC_BASE + 4*acc_idx : +3]. - # The scale operands are MFMA-ready packed dwords. mi/ni choose - # which of the four bytes inside the A/B scale dword the MFMA uses. - acc_pin = PIN_ACC_BASE + acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - _to_raw_inline_asm_operand(a_frag), - _to_raw_inline_asm_operand(b_frag), - _to_raw_inline_asm_operand(a_scale), - _to_raw_inline_asm_operand(b_scale), - ], - "v_mfma_scale_f32_16x16x128_f8f6f4 " - f"a[{acc_pin}:{acc_pin + 3}], " - "$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}], " - "$2, $3 " - f"op_sel:[{mi & 1},{ni & 1},0] " - f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", - has_side_effects=True, - ) + def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Fixed physical accumulator bank, visible SSA A/B/scale operands. + # acc_idx maps directly to a[PIN_ACC_BASE + 4*acc_idx : +3]. + # The scale operands are MFMA-ready packed dwords. mi/ni choose + # which of the four bytes inside the A/B scale dword the MFMA uses. + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(b_scale), + ], + "v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}], " + "$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", + has_side_effects=True, + ) - def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): - # Final-page form used by HK: destination and previous partial sum - # may be different AGPR ranges. Once old_acc_idx is consumed, its - # physical slot is dead and can be reused as a later destination. - dst_pin = PIN_ACC_BASE + dst_slot * 4 - old_pin = PIN_ACC_BASE + old_acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - _to_raw_inline_asm_operand(a_frag), - _to_raw_inline_asm_operand(b_frag), - _to_raw_inline_asm_operand(a_scale), - _to_raw_inline_asm_operand(b_scale), - ], - "v_mfma_scale_f32_16x16x128_f8f6f4 " - f"a[{dst_pin}:{dst_pin + 3}], " - "$0, $1, " - f"a[{old_pin}:{old_pin + 3}], " - "$2, $3 " - f"op_sel:[{mi & 1},{ni & 1},0] " - f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", - has_side_effects=True, - ) + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Final-page form used by HK: destination and previous partial sum + # may be different AGPR ranges. Once old_acc_idx is consumed, its + # physical slot is dead and can be reused as a later destination. + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(b_scale), + ], + "v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + "$0, $1, " + f"a[{old_pin}:{old_pin + 3}], " + "$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", + has_side_effects=True, + ) - def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): - """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" - mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE - pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, 0) - pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, 1) - pinned_mfma(acc_base + 2, a_frag, b2, a_scale, bs2, mi, 2) - pinned_mfma(acc_base + 3, a_frag, b3, a_scale, bs3, mi, 3) - - def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): - mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE - pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) - pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) - - def store_acc_vector_for_logical_idx(logical_acc_idx, acc): - subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - - # Bias depends only on the output-feature (N) coordinate, so read it - # once per column and reuse across the four M rows below. Index by - # the GLOBAL feature (by_n_idx + col), matching the C store's - # c_tile_base_elems fold; the tile-local col alone would reread - # bias[0..BLOCK_N) for every N-block. - if const_expr(has_bias): - bias_value = load_bias(by_n_idx + col) + def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): + """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, 1) + pinned_mfma(acc_base + 2, a_frag, b2, a_scale, bs2, mi, 2) + pinned_mfma(acc_base + 3, a_frag, b3, a_scale, bs3, mi, 3) + + def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 - # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing: - # value = acc [+ bias] (pre-activation) - # GELU_AUX: save pre-activation to Aux, then value = gelu(value) - value = Vec(acc)[ii] + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column and reuse across the four M rows below. Index by + # the GLOBAL feature (by_n_idx + col), matching the C store's + # c_tile_base_elems fold; the tile-local col alone would reread + # bias[0..BLOCK_N) for every N-block. if const_expr(has_bias): - value = value + bias_value + bias_value = load_bias(by_n_idx + col) + + for ii in range_constexpr(4): + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) + value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + + if const_expr(has_gelu): + # Save the pre-activation (post-bias) value for backward, + # then apply GELU to the C output. + aux_val = value + if output_dtype != torch.float32: + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) - if const_expr(has_gelu): - # Save the pre-activation (post-bias) value for backward, - # then apply GELU to the C output. - aux_val = value if output_dtype != torch.float32: - aux_val = aux_val.to(output_fx_dtype) - aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) - fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) - value = gelu_tanh(value) - - if output_dtype != torch.float32: - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) - - # Explicit register coordinates for HK-style four-quadrant mapping. - # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions - # inside each 128x128 quadrant: - # cA: (warp_m, warp_n) - # cB: (warp_m, warp_n + 2) - # cC: (warp_m + 2, warp_n) - # cD: (warp_m + 2, warp_n + 2) - reg_lds_k_col0, reg_lds_k_col1 = _normal_read_columns( - lane_div_16, - lane_mod_16, - ) - - reg_subtile_m_idx0 = wave_id // 2 - reg_subtile_n_idx0 = wave_id % 2 - - reserve_pinned_accumulators() - zero_pinned_accumulators() - - def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_scales = scale_tile[2] if sn == 0 else scale_tile[3] - - b_ni = _load_b_ni( - load_transposed_frag, - load_normal_b_frag, - lds_b, - sn, - ni, - reg_subtile_n_idx0, - lane_mod_16, - ) - return b_ni, b_scales[ni] - - def load_b_subtile_regs(lds_b, scale_tile, sn): - b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) - b1, bs1 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 1) - b2, bs2 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 2) - b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) - return b0, b1, b2, b3, bs0, bs1, bs2, bs3 - - def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - - return _load_a_half( - load_transposed_frag_half, - load_frag_half_at_byte_base, - lds_a, - sm, - mi, - half, - reg_subtile_m_idx0, + value = value.to(output_fx_dtype) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_lds_k_col0, reg_lds_k_col1 = _normal_read_columns( + lane_div_16, lane_mod_16, ) - def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): - # Fine-grained A register load for one 16-row M-direction MFMA slice. - a_scales = scale_tile[0] if sm == 0 else scale_tile[1] - x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) - x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) - a_mi = pack_frag_halves(x0, x1) - a_scale_mi = a_scales[mi] - return a_mi, a_scale_mi - - def load_a_subtile_regs(lds_a, scale_tile, sm): - a0, as0 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 0) - a1, as1 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 1) - a2, as2 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 2) - a3, as3 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 3) - return a0, a1, a2, a3, as0, as1, as2, as3 - - def hk_one_k_with_refill( - k128, - cur_a, - cur_b, - next_a, - next_b, - refill_a, - refill_b, - a0_regs, - b0_regs, - cur_scales, - prev_refill_scales, - ): - # Scale invariant: - # cur_scales is HK MFMA-ready for K. - # prev_refill_scales is HK MFMA-ready for K+1. - # This iteration issues K+2 scale loads and returns them for the - # next steady iteration or final tail. - - # Wait only far enough for the current page; the next-page refill may remain in flight. - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 - # Immediately issue MFMA-ready K+2 scale loads. - # They are returned for the next iteration without any in-kernel - # byte extraction or broadcast. - refill_scales = load_scale_tile(fx.Index(k128 + 2)) - next_scales_ready = prev_refill_scales - # A-top and B-left are both carried as complete 64-row register tiles, - # so their LDS half-pages can be refilled immediately. - a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs - b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - - b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) - b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) - b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) - b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) - - # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) - - # Q0: interleave the current tile's A-bottom LDS reads with K+2 - # refills and Q0 compute. Each complete A-bottom fragment is assembled - # from two independently scheduled K64 halves. - rocdl.sched_barrier(0) - a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a) - mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) + reserve_pinned_accumulators() + zero_pinned_accumulators() - a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) - stage_b_subtile_pass(k_refill, 0, 0, refill_b) - mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) + def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_scales = scale_tile[2] if sn == 0 else scale_tile[3] - a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) - stage_a_subtile_pass(k_refill, 0, 1, refill_a) - mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) + b_ni = _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) + return b_ni, b_scales[ni] + + def load_b_subtile_regs(lds_b, scale_tile, sn): + b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) + b1, bs1 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 1) + b2, bs2 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 2) + b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) + return b0, b1, b2, b3, bs0, bs1, bs2, bs3 + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) - a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) - stage_b_subtile_pass(k_refill, 0, 1, refill_b) - mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) + def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): + # Fine-grained A register load for one 16-row M-direction MFMA slice. + a_scales = scale_tile[0] if sm == 0 else scale_tile[1] + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + a_mi = pack_frag_halves(x0, x1) + a_scale_mi = a_scales[mi] + return a_mi, a_scale_mi + + def load_a_subtile_regs(lds_a, scale_tile, sm): + a0, as0 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 0) + a1, as1 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 1) + a2, as2 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 2) + a3, as3 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 3) + return a0, a1, a2, a3, as0, as1, as2, as3 + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + cur_scales, + prev_refill_scales, + ): + # Scale invariant: + # cur_scales is HK MFMA-ready for K. + # prev_refill_scales is HK MFMA-ready for K+1. + # This iteration issues K+2 scale loads and returns them for the + # next steady iteration or final tail. + + # Wait only far enough for the current page; the next-page refill may remain in flight. + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # Immediately issue MFMA-ready K+2 scale loads. + # They are returned for the next iteration without any in-kernel + # byte extraction or broadcast. + refill_scales = load_scale_tile(fx.Int64(k128 + 2)) + next_scales_ready = prev_refill_scales + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Int64((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + as10 = cur_scales[1][0] + as11 = cur_scales[1][1] + as12 = cur_scales[1][2] + as13 = cur_scales[1][3] + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + barrier( + vmcnt=2 * LOAD_PASSES_A_SUBTILE + + 2 * LOAD_PASSES_B_SUBTILE + + LOAD_PASSES_SCALES, + lgkmcnt=0, + ) + rocdl.sched_barrier(0) - a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) - stage_a_subtile_pass(k_refill, 0, 2, refill_a) - mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) - stage_b_subtile_pass(k_refill, 0, 2, refill_b) - mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) - stage_a_subtile_pass(k_refill, 0, 3, refill_a) - mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) - stage_b_subtile_pass(k_refill, 0, 3, refill_b) - mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - hot_loop_scheduler_q0_refill_a1_2n() + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - # Retire the eight distributed A-bottom LDS reads before K+2 refills - # overwrite the current page's A-bottom half-page. Keep this wait as - # late as possible to maximize read/compute overlap. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - a10 = pack_frag_halves(a10_x0, a10_x1) - a11 = pack_frag_halves(a11_x0, a11_x1) - a12 = pack_frag_halves(a12_x0, a12_x1) - a13 = pack_frag_halves(a13_x0, a13_x1) - as10 = cur_scales[1][0] - as11 = cur_scales[1][1] - as12 = cur_scales[1][2] - as13 = cur_scales[1][3] + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 0, refill_b) - mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - stage_a_subtile_pass(k_refill, 1, 0, refill_a) - mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + hot_loop_scheduler_q_prefetch_4n() - stage_b_subtile_pass(k_refill, 1, 1, refill_b) - mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) - stage_a_subtile_pass(k_refill, 1, 1, refill_a) - mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales - stage_b_subtile_pass(k_refill, 1, 2, refill_b) - mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) + def hk_one_k_tail_with_next( + cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales + ): + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - stage_a_subtile_pass(k_refill, 1, 2, refill_a) - mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - stage_b_subtile_pass(k_refill, 1, 3, refill_b) - mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) - stage_a_subtile_pass(k_refill, 1, 3, refill_a) - mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) - hot_loop_scheduler_q_refill_2n() + mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - # Leave exactly the K+2 refill and scale loads outstanding. The following - # LDS reads consume the already-ready next page, not the page being refilled. - rocdl.sched_barrier(0) - barrier( - vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, - lgkmcnt=0, - ) - rocdl.sched_barrier(0) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) - next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + rocdl.sched_barrier(0) + barrier(vmcnt=LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - hot_loop_scheduler_q_prefetch_4n() + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - next_a0_regs = ( - next_a00, - next_a01, - next_a02, - next_a03, - next_as00, - next_as01, - next_as02, - next_as03, - ) - next_b0_regs = ( - next_b00, - next_b01, - next_b02, - next_b03, - next_bs00, - next_bs01, - next_bs02, - next_bs03, - ) + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - def hk_one_k_tail_with_next( - cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales - ): - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + hot_loop_scheduler_q_prefetch_4n() - a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs - b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) - b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) - b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) - b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) - b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): + barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + a_scales = (as00, as01, as02, as03, as10, as11, as12, as13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + b_scales = (bs00, bs01, bs02, bs03, bs10, bs11, bs12, bs13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + a_scales[a_frag_idx], + b_scales[b_frag_idx], + mi, + ni, + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) - mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in + # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], + # and load_scale_tile returns the current wave's scale operands in VGPRs. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + # Load scales first, so that they become the oldest VMEM ops. + scales0 = load_scale_tile(fx.Int64(0)) + scales1 = load_scale_tile(fx.Int64(1)) - a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) - a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) - a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) - a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) - mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) - barrier(vmcnt=LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) rocdl.sched_barrier(0) - next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + # scales0 is already MFMA-ready; no byte extraction or broadcast is needed. + # Keep the hot loop consistent for k=0 and k>0: + # K0 is consumed directly. K1 MFMA-ready scales are carried as + # prev_refill_scales and become next_scales_ready at loop entry. - next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - hot_loop_scheduler_q_prefetch_4n() - - next_a0_regs = ( - next_a00, - next_a01, - next_a02, - next_a03, - next_as00, - next_as01, - next_as02, - next_as03, - ) - next_b0_regs = ( - next_b00, - next_b01, - next_b02, - next_b03, - next_bs00, - next_bs01, - next_bs02, - next_bs03, - ) - - return next_a0_regs, next_b0_regs - - def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): - barrier(vmcnt=0, lgkmcnt=0) - - a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs - b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - - # Materialize the remaining final-page A/B fragments once. The - # subsequent schedule is entirely register/AGPR traffic. - b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) - b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) - b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) - b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + # Seed the carried-register pipeline with K0 A-top. In later steady-state + # iterations, Q2/Q3 of the preceding iteration prefetch the next tile's + # A-top and B-left register tiles before their LDS half-pages are reused. + a0_regs = load_a_subtile_regs(lds_a0, scales0, 0) rocdl.sched_barrier(0) - barrier(lgkmcnt=0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) rocdl.sched_barrier(0) - a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) - a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) - a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) - a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) - - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) - - a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) - a_scales = (as00, as01, as02, as03, as10, as11, as12, as13) - b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) - b_scales = (bs00, bs01, bs02, bs03, bs10, bs11, bs12, bs13) - - # Rolling final-page epilogue. - # - # Finalize accumulators in their own physical AGPR slots, but delay - # each AGPR read/store until several independent final MFMAs have - # been issued. - # - # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, - # MFMA 4, drain 1, MFMA 5, drain 2, ... - # - # The buffer stores are only issued here; they may remain in flight - # while later MFMAs and accumulator drains continue. - FINAL_EPILOGUE_DEPTH = 4 - pending = [] - - for old_acc_idx in range_constexpr(ACCS_PER_WAVE): - subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi - b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni - - # Final MFMA remains in-place. The logical accumulator's own - # AGPR slot is unique and cannot conflict with another pending - # result, so no ad-hoc physical-slot permutation is needed. - pinned_final_mfma( - old_acc_idx, - old_acc_idx, - a_frags[a_frag_idx], - b_frags[b_frag_idx], - a_scales[a_frag_idx], - b_scales[b_frag_idx], - mi, - ni, - ) - pending.append(old_acc_idx) - - # Drain the oldest completed result only after enough newer - # independent MFMAs have supplied the MFMA->AGPR-read spacing. - if len(pending) == FINAL_EPILOGUE_DEPTH: - drain_acc_idx = pending.pop(0) - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) - - # Flush the final results after all final-page MFMAs have issued. - for drain_acc_idx in pending: - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) - - # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in - # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], - # and load_scale_tile returns the current wave's scale operands in VGPRs. - - # Load scales first, so that they become the oldest VMEM ops. - scales0 = load_scale_tile(fx.Index(0)) - scales1 = load_scale_tile(fx.Index(1)) - - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) - - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) - - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) - - # scales0 is already MFMA-ready; no byte extraction or broadcast is needed. - # Keep the hot loop consistent for k=0 and k>0: - # K0 is consumed directly. K1 MFMA-ready scales are carried as - # prev_refill_scales and become next_scales_ready at loop entry. - - # Seed the carried-register pipeline with K0 A-top. In later steady-state - # iterations, Q2/Q3 of the preceding iteration prefetch the next tile's - # A-top and B-left register tiles before their LDS half-pages are reused. - a0_regs = load_a_subtile_regs(lds_a0, scales0, 0) - - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) - - # Complete the K0 carried-register seed with B-left. - b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) - - # Main HK loop: exactly one logical K128 per iteration. - # Even k consumes and refills LDS0; odd k does the same for LDS1. - # Scale tiles follow the same K128 progression but remain in VGPRs. - refill_scales = scales1 # K1 scales become the next ready scale tile at loop entry - for k128 in range_constexpr(NUM_K_TILES - 2): - if (k128 % 2) == 0: - a0_regs, b0_regs, scales1, refill_scales = hk_one_k_with_refill( - k128, + # Complete the K0 carried-register seed with B-left. + b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + # Scale tiles follow the same K128 progression but remain in VGPRs. + refill_scales = scales1 # K1 scales become the next ready scale tile at loop entry + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs, scales1, refill_scales = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + scales0, + refill_scales, + ) + else: + a0_regs, b0_regs, scales0, refill_scales = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + scales1, + refill_scales, + ) + + # Common two-page tail. The penultimate tile still uses the Q2/Q3 + # carry-prefetch scheduler to prepare A-top/B-left for the final tile, + # but it performs no K+2 data or scale refill. The final tile performs + # compute only. After the steady loop, a0_regs/b0_regs belong to the + # next tile to consume, while refill_scales belongs to the page most + # recently refilled; therefore tail page order depends on parity: + # even NUM_K_TILES: consume LDS0 then final LDS1 + # odd NUM_K_TILES: consume LDS1 then final LDS0 + if (NUM_K_TILES % 2) == 0: + scales1 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a0, lds_b0, lds_a1, lds_b1, - lds_a0, - lds_b0, a0_regs, b0_regs, scales0, - refill_scales, + scales1, ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) else: - a0_regs, b0_regs, scales0, refill_scales = hk_one_k_with_refill( - k128, + scales0 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a1, lds_b1, lds_a0, lds_b0, - lds_a1, - lds_b1, a0_regs, b0_regs, scales1, - refill_scales, + scales0, ) - - # Common two-page tail. The penultimate tile still uses the Q2/Q3 - # carry-prefetch scheduler to prepare A-top/B-left for the final tile, - # but it performs no K+2 data or scale refill. The final tile performs - # compute only. After the steady loop, a0_regs/b0_regs belong to the - # next tile to consume, while refill_scales belongs to the page most - # recently refilled; therefore tail page order depends on parity: - # even NUM_K_TILES: consume LDS0 then final LDS1 - # odd NUM_K_TILES: consume LDS1 then final LDS0 - if (NUM_K_TILES % 2) == 0: - scales1 = refill_scales - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a0, - lds_b0, - lds_a1, - lds_b1, - a0_regs, - b0_regs, - scales0, - scales1, - ) - hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) + + # Dispatch. Dense runs one XCD-swizzled tile. Grouped runs a persistent + # loop over precomputed (pid_m, pid_n, expert) descriptors, selecting the + # schedule at compile time (``sched``): + # "worksteal": each iteration claims the next global slot via one atomic + # fetch-add (dynamic load balancing across CUs). + # "static": each workgroup owns a fixed grid-stride slice of slots + # (slot = blockIdx.x, blockIdx.x + gridDim.x, ...). No atomic + # or LDS broadcast, so it is deadlock-proof; used as a + # heuristic fast path and a correctness reference. + def _run_grouped_tile(slot): + pid_m = _desc_field(slot, fx.Int32(0)) + pid_n = _desc_field(slot, fx.Int32(1)) + expert = _desc_field(slot, fx.Int32(2)) + process_tile(pid_m, pid_n, expert) + + if const_expr(grouped): + if const_expr(sched == "static"): + stride = fx.Int32(gpu.grid_dim.x) + slot = fx.Int32(gpu.block_id("x")) + while slot < num_tiles: + _run_grouped_tile(slot) + slot = slot + stride + else: + # Single persistent loop, mirroring mega_moe_stage1: claim the next + # slot at the TOP of the loop body (atomic + barrier + LDS broadcast, + # so every thread sees the same slot), run the tile only when the + # slot is in range, and derive the loop-continue flag from that same + # broadcast slot. Keeping the barrier at one site with a uniform + # continue condition guarantees convergent barriers. + active = fx.Int32(1) == fx.Int32(1) + while active: + slot = _claim_next_slot() + in_range = slot < num_tiles + if in_range: + _run_grouped_tile(slot) + active = in_range else: - scales0 = refill_scales - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a1, - lds_b1, - lds_a0, - lds_b0, - a0_regs, - b0_regs, - scales1, - scales0, - ) - hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) + process_tile(pid_m, pid_n, fx.Int32(0)) @flyc.jit def launch_gemm( @@ -1500,10 +1875,21 @@ def launch_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, stream: fx.Stream = fx.Stream(None), ): # The integration only dispatches aligned shapes; no partial-tile masking exists. + # Dense launch: one workgroup per output tile. Grouped launch: a fixed + # num-CU grid of persistent workgroups that steal tiles from the atomic + # work queue (num_tiles spans the packed [M_total, N] output). grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(group or wgrad): + grid_x = get_num_cus() kernel_gemm( A, As, @@ -1514,6 +1900,12 @@ def launch_gemm( Aux, c_m, c_n, + num_groups, + work_queue, + work_counter, + b_data_ptrs, + a_data_ptrs, + num_work, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1532,25 +1924,55 @@ def do_gemm( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, + num_work: int = 0, + a_shape=None, + b_shape=None, + sched: str = "worksteal", + wgrad: bool = False, ): """Launch one cached compile-time MXFP8 layout specialization. The GELU epilogues require a caller-allocated ``aux`` output (M x N, same dtype as C), filled in place with the pre-activation values for backward. + + With ``a_data_ptrs`` the per-expert A operand is read in place (see act_ptrs in + ``_compile_kernel``): ``A`` is then one representative expert and ``a_shape`` the + packed operand shape. ``b_shape`` likewise gives wgrad's packed B shape when B + is read through ``b_data_ptrs``. """ + a_rows, a_cols = a_shape if a_shape is not None else A.shape + b_rows, b_cols = b_shape if b_shape is not None else B.shape if layout == "TN": - M_runtime, K_runtime = A.shape - N_runtime, Kb_runtime = B.shape + M_runtime, K_runtime = a_rows, a_cols + N_runtime, Kb_runtime = b_rows, b_cols elif layout == "NN": - M_runtime, K_runtime = A.shape - Kb_runtime, N_runtime = B.shape + M_runtime, K_runtime = a_rows, a_cols + Kb_runtime, N_runtime = b_rows, b_cols elif layout == "NT": - K_runtime, M_runtime = A.shape - Kb_runtime, N_runtime = B.shape + K_runtime, M_runtime = a_rows, a_cols + Kb_runtime, N_runtime = b_rows, b_cols else: raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + # Grouped wgrad (NT): dY and X stack the G experts on the contraction (token) + # axis as G*M_pad, while the kernel's compile-time contraction is the per-expert + # M_pad. M_runtime (N_out) and N_runtime (K_in) are already per-expert. + if wgrad: + assert ( + K_runtime % num_groups == 0 + ), f"grouped wgrad contraction {K_runtime} not divisible by num_groups={num_groups}" + K_runtime = K_runtime // num_groups + + # Grouped fwd/dgrad: N is the per-expert output width, C's last dim (TN: N_out, + # NN: K_in). wgrad keeps M/N per-expert; its C stacks experts on the M axis. + if group: + N_runtime = C.shape[-1] supported_fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" @@ -1564,25 +1986,43 @@ def do_gemm( block_k=_BLOCK_K, label=f"MXFP8 {layout} GEMM", ) + # With per-expert pointers, A (and wgrad's B) is one representative expert, but + # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit + # offsets, so the bound is on the packed extent. require_launch_size( f"MXFP8 {layout} GEMM", - ("A", A), - ("B", B), + ("A", A, a_rows * a_cols * A.element_size()), + ("B", B, b_rows * b_cols * B.element_size()), ("As", As), ("Bs", Bs), ("C", C), ) - expected_as = (K_runtime // _BLOCK_K, M_runtime) - expected_bs = (K_runtime // _BLOCK_K, N_runtime) + # Packed scale shapes. + if wgrad: + # wgrad (NT): both A (dY) and B (X) scales are packed on the contraction + # axis to G*M_pad/128. A-scale [G*M_pad/128, N_out], B-scale [.., K_in]. + packed_k128 = (K_runtime * num_groups) // _BLOCK_K + expected_as = (packed_k128, M_runtime) + expected_bs = (packed_k128, N_runtime) + else: + expected_as = (K_runtime // _BLOCK_K, M_runtime) + # TN packs the B-scale along N (column count = G*N_out); NN packs it along the + # contraction axis (row count on K128 = G*N_out/128). Kb_runtime is the + # per-expert N_out, since the B data is addressed per expert. + if group and layout == "NN": + expected_bs = ((Kb_runtime * num_groups) // _BLOCK_K, N_runtime) + else: + n_packed = N_runtime * num_groups if group else N_runtime + expected_bs = (K_runtime // _BLOCK_K, n_packed) assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" assert tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" assert tuple(Bs.shape) == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" - assert tuple(C.shape) == ( - M_runtime, - N_runtime, - ), f"C shape {tuple(C.shape)} != {(M_runtime, N_runtime)}" + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad C is the + # unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + assert tuple(C.shape) == expected_c, f"C shape {tuple(C.shape)} != {expected_c}" if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" @@ -1624,7 +2064,7 @@ def do_gemm( raise ValueError(f"MXFP8 epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = torch.cuda.current_stream() + stream = grouped_utils.current_stream_handle(A.device) # A non-contiguous output cannot receive the kernel's writes: the flat # descriptor below would target a throwaway ``.contiguous()`` copy. Reject @@ -1632,30 +2072,75 @@ def do_gemm( if not C.is_contiguous(): raise FlyDSLUnsupportedError("FlyDSL MXFP8 requires contiguous output storage") - # Preserve the exact flat descriptor contract used by the passing kernels. + # The kernel addresses every operand through a flat 1-D buffer descriptor. A_arg = A.view(torch.uint8).contiguous().view(-1) B_arg = B.view(torch.uint8).contiguous().view(-1) As_arg = As.contiguous().view(-1) Bs_arg = Bs.contiguous().view(-1) C_arg = C.view(-1) - # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. These + # placeholders are never dereferenced by the kernel, so a cached shared buffer + # (see ``placeholder_arg``) is used rather than a fresh allocate+zero every call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + Aux_arg = placeholder_arg(A.device, C.dtype) + + # Grouped runtime tensors are required by the kernel signature; the dense + # compiled kernel never dereferences them but still needs placeholder + # buffers. wgrad uses the same persistent descriptor queue as the grouped fwd/dgrad path. + if group or wgrad: + # num_work is the kernel's drain bound: it must be the queue's descriptor count. + if work_queue.numel() != 3 * int(num_work): + raise ValueError( + f"num_work={num_work} does not match work_queue of {work_queue.numel()} " + "int32 (3 per descriptor)" + ) + work_queue_arg = work_queue.contiguous().view(-1) + work_counter_arg = work_counter.contiguous().view(-1) + else: + work_queue_arg = placeholder_arg(A.device, torch.int32) + work_counter_arg = placeholder_arg(A.device, torch.int32) + + # Per-expert B base pointers: the weights in grouped fwd/dgrad, or wgrad's + # pre-shifted X with act_ptrs. Dense passes an unused placeholder so the + # positional kernel signature stays the same. + if b_data_ptrs is not None: + b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) + else: + b_data_ptrs_arg = placeholder_arg(A.device, torch.int64) + # act_ptrs also selects how the kernel reads B: outside grouped fwd/dgrad, B is + # read through b_data_ptrs only when act_ptrs is set. So activation pointers are + # grouped-only, and wgrad takes the dY and X pointer arrays together or not at all. + if a_data_ptrs is not None and not (group or wgrad): + raise ValueError("MXFP8: per-expert activation pointers require a grouped launch") + if wgrad and (a_data_ptrs is None) != (b_data_ptrs is None): + raise ValueError("MXFP8 wgrad: a_data_ptrs and b_data_ptrs must be given together") + # Per-expert activation base pointers (see act_ptrs in _compile_kernel). + act_ptrs = a_data_ptrs is not None + if act_ptrs: + a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) + else: + a_data_ptrs_arg = placeholder_arg(A.device, torch.int64) - _cached_launch( + launch_key = ( + "mxfp8_gemm", K_runtime, A.dtype, B.dtype, C.dtype, layout, epilogue, - )( + group, + sched, + wgrad, + act_ptrs, + ) + launch_args = ( A_arg, As_arg, B_arg, @@ -1665,8 +2150,21 @@ def do_gemm( Aux_arg, M_runtime, N_runtime, - stream=stream, + num_groups, + work_queue_arg, + work_counter_arg, + b_data_ptrs_arg, + a_data_ptrs_arg, + int(num_work), + stream, ) + compiled_launch( + launch_key, + _cached_launch( + K_runtime, A.dtype, B.dtype, C.dtype, layout, epilogue, group, sched, wgrad, act_ptrs + ), + launch_args, + )(*launch_args) @functools.lru_cache(maxsize=None) @@ -1677,8 +2175,12 @@ def _cached_launch( output_dtype: torch.dtype, layout: str, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, + act_ptrs: bool = False, ): - """Cache independent TN/NN/NT binaries with no runtime layout argument.""" + """Cache independent TN/NN/NT compiled kernels with no runtime layout argument.""" return _compile_kernel( K, a_fp8_dtype, @@ -1686,6 +2188,10 @@ def _cached_launch( output_dtype, layout, epilogue, + group, + sched, + wgrad, + act_ptrs, ) @@ -1715,6 +2221,74 @@ def _validate_common_payloads( ) +def _mxfp8_matmul_grouped_wgrad( + a, # dY packed [G*M_pad, N_out], columnwise fp8 + a_scale, # [G*M_pad/32, N_out] raw E8M0 uint8, or list[G] of per-expert scales + b, # X packed [G*M_pad, K_in], columnwise fp8 + b_scale, # [G*M_pad/32, K_in] raw E8M0 uint8, or list[G] of per-expert scales + D, # dW packed [G*N_out, K_in] + *, + stream, + num_groups, + work_queue, + work_counter, + sched, + a_data_ptrs=None, + b_data_ptrs=None, + num_work=0, + a_shape=None, + b_shape=None, +): + """Grouped wgrad (NT) launch: dW = dY^T @ X, contracting over padded tokens. + + Both operands transpose-read (columnwise) and contract over the token axis; the + kernel's compile-time contraction is the per-expert padded token count + M_pad = (G*M_pad)/G. Experts are packed on the contraction axis for a/b and on + the output-M axis for D. Operands arrive already packed + MXFP8-quantized, or + (``a_data_ptrs``/``b_data_ptrs``) as one representative expert each plus + per-expert pointers, with ``a_shape``/``b_shape`` the packed shapes. + """ + gm = (a_shape if a_shape is not None else a.shape)[0] # gm = G*M_pad + gm_b = (b_shape if b_shape is not None else b.shape)[0] + if gm != gm_b: + raise ValueError(f"wgrad a/b token dims differ: {gm} vs {gm_b}") + if gm % num_groups != 0: + raise ValueError(f"wgrad packed tokens {gm} not divisible by num_groups={num_groups}") + m_pad = gm // num_groups + if m_pad % SCALE_GROUP_SIZE != 0: + raise FlyDSLUnsupportedError( + f"wgrad padded per-expert tokens M_pad={m_pad} must be divisible by {SCALE_GROUP_SIZE}" + ) + if _scale_dtype(a_scale) != torch.uint8 or _scale_dtype(b_scale) != torch.uint8: + raise TypeError("FlyDSL MXFP8 wgrad expects raw E8M0 scales as torch.uint8") + + # NT transpose-read: both scales are packed from their columnwise (K-major) + # source into the kernel's [K/128, dim] int32 layout. + a_scale_hk = _pack_scales(a_scale, source_colwise=True, stream=stream) + b_scale_hk = _pack_scales(b_scale, source_colwise=True, stream=stream) + + do_gemm( + a, + a_scale_hk, + b, + b_scale_hk, + D, + layout="NT", + stream=stream, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, + a_shape=a_shape, + b_shape=b_shape, + num_work=num_work, + sched=sched, + wgrad=True, + ) + return D + + def mxfp8_matmul( a: torch.Tensor, a_scale: torch.Tensor, @@ -1727,40 +2301,93 @@ def mxfp8_matmul( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, + num_work: int = 0, + a_shape=None, + b_shape=None, + sched: str = "worksteal", + wgrad: bool = False, ): - """Normalize scale orientation and launch a compile-time layout binary. + """Normalize scale orientation and launch a kernel compiled for a fixed layout. Wrapper-visible contracts: TN: a [M,K], b [N,K], scales [M,K/32] and [N,K/32] + (group): a [M_total,K], b [G*N,K] scales [M_total,K/32] and [G*N,K/32] NN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] + (group): a [M_total,Kc], b [G*Kc,N] scales [M_total,Kc/32] and + [G*Kc/32,N] -- grouped NN is the dgrad dX=dY@W: the kernel + contraction Kc is N_out and experts pack on that axis. NT: a [K,M], b [K,N], scales [K/32,M] and [K/32,N] TN consumes the selected rowwise payloads directly. NN and NT preserve K-major payloads and use ``ds_read_b64_tr_b8`` inside their compile-time-specialized kernels. + + ``a_data_ptrs`` reads the grouped A operand per expert in place: ``a`` is then one + representative expert and ``a_shape`` the packed shape (see ``do_gemm``). A scale + may be a list of per-expert tensors stacked on dim 0, packed in place + (:func:`pack_mx32_scales_for_hk_experts`). """ + + if wgrad: + return _mxfp8_matmul_grouped_wgrad( + a, + a_scale, + b, + b_scale, + D, + stream=stream, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + a_data_ptrs=a_data_ptrs, + b_data_ptrs=b_data_ptrs, + a_shape=a_shape, + b_shape=b_shape, + num_work=num_work, + ) + + if group and layout not in ("TN", "NN"): + raise ValueError( + "FlyDSL MXFP8 grouped fwd/dgrad supports only the TN and NN layouts " + "(grouped NT wgrad uses wgrad=True)" + ) + if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") _validate_common_payloads(a, b, D, layout=layout) + a_rows, a_cols = a_shape if a_shape is not None else a.shape if layout == "TN": - m, k = a.shape + m, k = a_rows, a_cols n, kb = b.shape elif layout == "NN": - m, k = a.shape + m, k = a_rows, a_cols kb, n = b.shape else: - k, m = a.shape + k, m = a_rows, a_cols kb, n = b.shape + # Grouped fwd (TN) and dgrad (NN) read B per expert through b_data_ptrs, so b is + # one expert's weight: its shape gives the per-expert output width (TN: N_out, + # NN: K_in) and contraction, exactly as in the dense case. Only the B scales are + # stacked over the G experts (checked below). + if group and b_data_ptrs is None: + raise FlyDSLUnsupportedError("grouped MXFP8 GEMM requires per-expert weight pointers") + n_out = n if kb != k: raise ValueError( f"Incompatible MXFP8 {layout} operands: A{tuple(a.shape)} and B{tuple(b.shape)}" ) - if tuple(D.shape) != (m, n): - raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n)}") + if k % SCALE_GROUP_SIZE != 0: raise FlyDSLUnsupportedError( f"K={k} must be divisible by MXFP8 scale group size {SCALE_GROUP_SIZE}" @@ -1772,21 +2399,26 @@ def mxfp8_matmul( expected_a_scale = (m, k // SCALE_GROUP_SIZE) if layout == "TN": - expected_b_scale = (n, k // SCALE_GROUP_SIZE) + # Grouped TN packs the B scale along N, spanning G*N_out rows (n == N_out). + b_scale_rows = n_out * num_groups if group else n + expected_b_scale = (b_scale_rows, k // SCALE_GROUP_SIZE) else: - expected_b_scale = (k // SCALE_GROUP_SIZE, n) + # NN/NT B scale is [contraction/32, N]. Grouped NN packs the scale along the + # contraction axis: kb is the per-expert N_out, so it spans kb*num_groups rows. + kb_packed = kb * num_groups if group else kb + expected_b_scale = (kb_packed // SCALE_GROUP_SIZE, n) - if tuple(a_scale.shape) != expected_a_scale: + if _scale_shape(a_scale) != expected_a_scale: raise ValueError( - f"a_scale shape {tuple(a_scale.shape)} != expected {expected_a_scale} for {layout}" + f"a_scale shape {_scale_shape(a_scale)} != expected {expected_a_scale} for {layout}" ) - if tuple(b_scale.shape) != expected_b_scale: + if _scale_shape(b_scale) != expected_b_scale: raise ValueError( - f"b_scale shape {tuple(b_scale.shape)} != expected {expected_b_scale} for {layout}" + f"b_scale shape {_scale_shape(b_scale)} != expected {expected_b_scale} for {layout}" ) - if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + if _scale_dtype(a_scale) != torch.uint8 or _scale_dtype(b_scale) != torch.uint8: raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") - if a_scale.device != a.device or b_scale.device != a.device: + if _scale_device(a_scale) != a.device or _scale_device(b_scale) != a.device: raise ValueError("A, B, scales, and D must be on the same device") if layout == "TN": @@ -1794,12 +2426,12 @@ def mxfp8_matmul( # a [M,K], b [N,K] a_kernel = a b_kernel = b - a_scale_hk = pack_mx32_scales_for_hk( + a_scale_hk = _pack_scales( a_scale, source_colwise=False, stream=stream, ) - b_scale_hk = pack_mx32_scales_for_hk( + b_scale_hk = _pack_scales( b_scale, source_colwise=False, stream=stream, @@ -1807,12 +2439,12 @@ def mxfp8_matmul( elif layout == "NN": a_kernel = a b_kernel = b - a_scale_hk = pack_mx32_scales_for_hk( + a_scale_hk = _pack_scales( a_scale, source_colwise=False, stream=stream, ) - b_scale_hk = pack_mx32_scales_for_hk( + b_scale_hk = _pack_scales( b_scale, source_colwise=True, stream=stream, @@ -1820,12 +2452,12 @@ def mxfp8_matmul( else: a_kernel = a b_kernel = b - a_scale_hk = pack_mx32_scales_for_hk( + a_scale_hk = _pack_scales( a_scale, source_colwise=True, stream=stream, ) - b_scale_hk = pack_mx32_scales_for_hk( + b_scale_hk = _pack_scales( b_scale, source_colwise=True, stream=stream, @@ -1843,12 +2475,21 @@ def mxfp8_matmul( a_scale_hk, b_kernel, b_scale_hk, - D.view(m, n), + D.view(m, n_out), layout=layout, stream=stream, epilogue=epilogue, bias=bias, - aux=aux.view(m, n) if aux is not None else None, + aux=aux.view(m, n_out) if aux is not None else None, + group=group, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, + a_shape=a_shape, + num_work=num_work, + sched=sched, ) return D diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1e039816d..da5f180d2 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -35,7 +35,7 @@ ) from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore from . import _split_quantization -from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..quantization import FP8GlobalStateManager, QuantizerRole, is_flydsl_grouped_gemm_enabled from ..utils import ( divide, cast_if_needed, @@ -749,6 +749,51 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( ) ) + @staticmethod + def _is_flydsl_grouped_gemm_supported( + *, + fp8, + recipe, + activation_dtype, + use_bias, + use_grouped_gemm_triton, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + in_features, + out_features, + ) -> bool: + """ROCm FlyDSL grouped GEMM: the opt-in and precision (see + ``is_flydsl_grouped_gemm_enabled``) plus this module's per-call constraints. + + It does not quantize: ``general_grouped_gemm_flydsl`` consumes TE's quantized + operands and still falls back per GEMM on anything its kernels cannot serve. + """ + if fp8: + enabled = recipe is not None and is_flydsl_grouped_gemm_enabled(recipe) + else: + # One half-precision kernel covers BF16 and FP16. + half_precision = activation_dtype in (torch.bfloat16, torch.float16) + enabled = half_precision and is_flydsl_grouped_gemm_enabled() + return ( + enabled + and in_features % 128 == 0 + and out_features % 128 == 0 + and not ( + use_bias + or use_grouped_gemm_triton + or backward_override + or cpu_offloading + or save_original_input + or debug + or unpad_output + or actual_m_splits is not None + ) + ) + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -1161,6 +1206,24 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) + # FlyDSL grouped GEMM (ROCm gfx950) opt-in: a drop-in for general_grouped_gemm + # over TE's already-quantized operands. + use_grouped_gemm_flydsl = _GroupedLinear._is_flydsl_grouped_gemm_supported( + fp8=fp8, + recipe=blockwise_recipe, + activation_dtype=activation_dtype, + use_bias=use_bias, + use_grouped_gemm_triton=use_grouped_gemm_triton, + backward_override=backward_override, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + in_features=weights[0].size(-1), + out_features=weights[0].size(0), + ) + backward_needs_input = is_grad_enabled and weight_requires_grad if backward_override is None and save_original_input and backward_needs_input: if recipe is not None and recipe.delayed(): @@ -1385,6 +1448,11 @@ def forward( if use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": m_splits_tensor} + elif use_grouped_gemm_flydsl: + from ..flydsl_kernels.gemm import general_grouped_gemm_flydsl + + general_grouped_gemm_func = general_grouped_gemm_flydsl + kwargs = {} else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} @@ -1522,6 +1590,7 @@ def forward( ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers ctx.use_grouped_gemm_triton = use_grouped_gemm_triton + ctx.use_grouped_gemm_flydsl = use_grouped_gemm_flydsl ctx.num_input_tensors = len(inputmats) ctx.dgrad_out = dgrad_out @@ -1990,6 +2059,11 @@ def backward( if ctx.use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + elif ctx.use_grouped_gemm_flydsl: + from ..flydsl_kernels.gemm import general_grouped_gemm_flydsl + + general_grouped_gemm_func = general_grouped_gemm_flydsl + kwargs = {} else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} @@ -2119,6 +2193,17 @@ def backward( if ctx.use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + elif ( + ctx.use_grouped_gemm_flydsl + and not ctx.fuse_wgrad_accumulation + and not (ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute()) + ): + # FlyDSL wgrad writes plain per-expert gradients; fused main_grad + # accumulation and delayed wgrad stay on the default grouped GEMM. + from ..flydsl_kernels.gemm import general_grouped_gemm_flydsl + + general_grouped_gemm_func = general_grouped_gemm_flydsl + kwargs = {} else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 176ee9e07..68a8085ad 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -331,6 +331,22 @@ def get_default_recipe() -> Recipe: return get_default_fp8_recipe() +def is_flydsl_grouped_gemm_enabled(recipe: Optional[Recipe] = None) -> bool: + """Whether the FlyDSL grouped GEMM (ROCm gfx950 opt-in) is enabled for ``recipe``. + + Covers the opt-in and the precisions the kernels implement: MXFP8 and tensor-wise + FP8 (current/delayed scaling) when ``recipe`` is given, else unquantized BF16/FP16. + GroupedLinear adds its per-call checks on top of this. + """ + if not ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" + and get_device_compute_capability() == (9, 5) + ): + return False + return recipe is None or recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed() + + def get_align_size_for_quantization(recipe: Recipe) -> int: """Get the alignment used to pad grouped quantized operations. @@ -343,6 +359,11 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: # alignment as the conservative fallback for context-free callers. if recipe.custom(): return recipe.quantization_alignment + # The FlyDSL grouped GEMM (gfx950 opt-in) tiles every expert with BLOCK_M=256, so + # per-expert token counts must be 256-aligned for the recipes it handles (MXFP8 and + # tensor-wise FP8); other recipes fall back and keep their own alignment. + if is_flydsl_grouped_gemm_enabled(recipe): + return 256 if recipe.mxfp8(): # HipKittens grouped GEMM requires 256-aligned expert dimensions. # HK is used by default when NVTE_USE_CUTLASS_GROUPED_GEMM=1,