From 06b9a89953a6340d07bfbfcc9ef2414effb13599 Mon Sep 17 00:00:00 2001 From: Mario Sieg Date: Wed, 26 Aug 2026 14:42:11 +0000 Subject: [PATCH] Cache the x casts to fp8 and build ragged layout for fp8 gemm once --- .../trainer/models/layers/fp8_grouped_gemm.py | 71 +++++++++++++++---- src/prime_rl/trainer/models/layers/moe.py | 28 ++++++-- 2 files changed, 79 insertions(+), 20 deletions(-) diff --git a/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py b/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py index bfc336029c..87fbc98f83 100644 --- a/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py +++ b/src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py @@ -102,7 +102,11 @@ def forward( x: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor, + layout: tuple | None = None, + x_fp8_cache: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: + if layout is None: + layout = build_grouped_layout(offs, total_m=x.size(0)) ( total_m, padded_total_m, @@ -112,19 +116,22 @@ def forward( starts_tensor, actual_ms_tensor, block_starts_tensor, - ) = build_grouped_layout(offs, total_m=x.size(0)) + ) = layout use_ue8m0 = ue8m0_for_device(x.device) - x_fp8 = grouped_per_token_cast_to_fp8_triton( - x, - padded_total_m, - block_to_group, - starts_tensor, - actual_ms_tensor, - block_starts_tensor, - use_ue8m0, - GROUP_ALIGNMENT, - ) + if x_fp8_cache is not None: + x_fp8 = x_fp8_cache + else: + x_fp8 = grouped_per_token_cast_to_fp8_triton( + x, + padded_total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + use_ue8m0, + GROUP_ALIGNMENT, + ) weight_fp8 = grouped_per_block_cast_to_fp8_triton( weight.transpose(1, 2), use_ue8m0, @@ -236,13 +243,44 @@ def backward(ctx, grad_output: torch.Tensor): block_starts_tensor, ) - return grad_x, grad_weight, None + # grad_x, grad_weight, offs, layout, x_fp8_cache — the last three never need grad. + return grad_x, grad_weight, None, None, None + + +def compute_grouped_layout(offs: torch.Tensor, total_m: int): + return build_grouped_layout(offs, total_m=total_m) + + +def cast_grouped_input_to_fp8(x: torch.Tensor, layout: tuple) -> tuple[torch.Tensor, torch.Tensor]: + ( + _total_m, + padded_total_m, + _grouped_layout, + block_to_group, + _ks_tensor, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + ) = layout + use_ue8m0 = ue8m0_for_device(x.device) + return grouped_per_token_cast_to_fp8_triton( + x, + padded_total_m, + block_to_group, + starts_tensor, + actual_ms_tensor, + block_starts_tensor, + use_ue8m0, + GROUP_ALIGNMENT, + ) def grouped_fp8_gemm( x: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor, + layout: tuple | None = None, + x_fp8_cache: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: """FP8 grouped GEMM, drop-in replacement for torch._grouped_mm. @@ -250,8 +288,15 @@ def grouped_fp8_gemm( x: (M, K) concatenated token activations in bfloat16. weight: (G, K, N) expert weights in bfloat16. offs: (G,) int32 cumulative token counts per expert. + layout: optional precomputed ``compute_grouped_layout(offs, total_m)`` result, + to skip rebuilding the ragged-group layout when the caller already has one + for this (offs, total_m) pair. + x_fp8_cache: optional precomputed ``cast_grouped_input_to_fp8(x, layout)`` + result, to skip re-casting ``x`` when the caller already cast this exact + tensor against this exact layout (e.g. for a sibling grouped_fp8_gemm call + with a different weight but the same x). Returns: (M, N) output tensor in bfloat16. """ - return _GroupedFP8Gemm.apply(x, weight, offs) + return _GroupedFP8Gemm.apply(x, weight, offs, layout, x_fp8_cache) diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 4f44cb5031..e4abb973ad 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -171,11 +171,23 @@ def _run_experts_grouped_mm_impl( assert x.dim() == 2 if fp8: - from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm + from prime_rl.trainer.models.layers.fp8_grouped_gemm import ( + cast_grouped_input_to_fp8, + compute_grouped_layout, + grouped_fp8_gemm, + ) + + x_bf16 = x.bfloat16() + layout = compute_grouped_layout(offsets, x_bf16.size(0)) + x_fp8_cache = cast_grouped_input_to_fp8(x_bf16, layout) - h = F.silu(grouped_fp8_gemm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offsets)) - h = h * grouped_fp8_gemm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offsets) - out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x) + h = F.silu( + grouped_fp8_gemm(x_bf16, w1.bfloat16().transpose(-2, -1), offsets, layout=layout, x_fp8_cache=x_fp8_cache) + ) + h = h * grouped_fp8_gemm( + x_bf16, w3.bfloat16().transpose(-2, -1), offsets, layout=layout, x_fp8_cache=x_fp8_cache + ) + out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets, layout=layout).type_as(x) else: h = F.silu(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets)) h = h * torch._grouped_mm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offs=offsets) @@ -1123,10 +1135,12 @@ def _run_nongated_experts_grouped_mm_impl( assert x.dim() == 2 if fp8: - from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm + from prime_rl.trainer.models.layers.fp8_grouped_gemm import compute_grouped_layout, grouped_fp8_gemm - h = relu2(grouped_fp8_gemm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offsets)) - out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x) + x_bf16 = x.bfloat16() + layout = compute_grouped_layout(offsets, x_bf16.size(0)) + h = relu2(grouped_fp8_gemm(x_bf16, w1.bfloat16().transpose(-2, -1), offsets, layout=layout)) + out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets, layout=layout).type_as(x) else: h = relu2(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets)) out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x)