|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +# DeepSpeed Team |
| 3 | +"""Fused AutoEP token restore without the eager scatter and FP32 intermediate. |
| 4 | +
|
| 5 | +The kernel reduces each token's top-k rows in FP32. Communication, routing, |
| 6 | +expert reorder, and grouped GEMM remain unchanged. |
| 7 | +""" |
| 8 | + |
| 9 | +from __future__ import annotations |
| 10 | + |
| 11 | +import torch |
| 12 | + |
| 13 | +from deepspeed.ops.triton_ops._triton import _TRITON_AVAILABLE, triton, tl |
| 14 | + |
| 15 | +_IS_ROCM_PYTORCH = getattr(torch.version, "hip", None) is not None |
| 16 | + |
| 17 | +SUPPORTED_ROW_DTYPES = (torch.bfloat16, torch.float16, torch.float32) |
| 18 | + |
| 19 | +_MAX_BLOCK_HIDDEN = 512 |
| 20 | +_INVERT_INDEX_BLOCK = 256 |
| 21 | +# The kernels hold a [slots, BLOCK_H] FP32 block live, so the hidden tile shrinks |
| 22 | +# as top-k grows to keep that block in registers rather than spilling. |
| 23 | +_MAX_BLOCK_ELEMENTS = 2048 |
| 24 | + |
| 25 | +if _TRITON_AVAILABLE: |
| 26 | + |
| 27 | + @triton.jit |
| 28 | + def _invert_index_kernel( |
| 29 | + index_ptr, |
| 30 | + inverse_ptr, |
| 31 | + num_indices, |
| 32 | + num_inverse_rows, |
| 33 | + BLOCK: tl.constexpr, |
| 34 | + ): |
| 35 | + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) |
| 36 | + in_range = offsets < num_indices |
| 37 | + |
| 38 | + targets = tl.load(index_ptr + offsets, mask=in_range, other=-1).to(tl.int64) |
| 39 | + writable = in_range & (targets >= 0) & (targets < num_inverse_rows) |
| 40 | + tl.store(inverse_ptr + tl.where(writable, targets, 0), offsets.to(tl.int32), mask=writable) |
| 41 | + |
| 42 | + @triton.jit |
| 43 | + def _weighted_restore_forward_kernel( |
| 44 | + rows_ptr, |
| 45 | + inverse_ptr, |
| 46 | + scores_ptr, |
| 47 | + out_ptr, |
| 48 | + hidden, |
| 49 | + rows_stride, |
| 50 | + scores_stride, |
| 51 | + out_stride, |
| 52 | + TOP_K: tl.constexpr, |
| 53 | + K_PADDED: tl.constexpr, |
| 54 | + BLOCK_H: tl.constexpr, |
| 55 | + ): |
| 56 | + token = tl.program_id(0) |
| 57 | + hidden_offsets = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) |
| 58 | + hidden_mask = hidden_offsets < hidden |
| 59 | + |
| 60 | + slots = tl.arange(0, K_PADDED) |
| 61 | + slot_mask = slots < TOP_K |
| 62 | + |
| 63 | + source_rows = tl.load(inverse_ptr + token * TOP_K + slots, mask=slot_mask, other=-1).to(tl.int64) |
| 64 | + row_valid = slot_mask & (source_rows >= 0) |
| 65 | + safe_rows = tl.where(row_valid, source_rows, 0) |
| 66 | + |
| 67 | + scores = tl.load(scores_ptr + token * scores_stride + slots, mask=slot_mask, other=0.0).to(tl.float32) |
| 68 | + |
| 69 | + block_mask = row_valid[:, None] & hidden_mask[None, :] |
| 70 | + values = tl.load( |
| 71 | + rows_ptr + safe_rows[:, None] * rows_stride + hidden_offsets[None, :], |
| 72 | + mask=block_mask, |
| 73 | + other=0.0, |
| 74 | + ).to(tl.float32) |
| 75 | + |
| 76 | + # Match the eager path's FP32 product and accumulation. |
| 77 | + weighted = tl.sum(values * scores[:, None], axis=0) |
| 78 | + tl.store( |
| 79 | + out_ptr + token * out_stride + hidden_offsets, |
| 80 | + weighted.to(out_ptr.dtype.element_ty), |
| 81 | + mask=hidden_mask, |
| 82 | + ) |
| 83 | + |
| 84 | + @triton.jit |
| 85 | + def _weighted_restore_backward_kernel( |
| 86 | + grad_out_ptr, |
| 87 | + rows_ptr, |
| 88 | + inverse_ptr, |
| 89 | + scores_ptr, |
| 90 | + grad_rows_ptr, |
| 91 | + grad_scores_ptr, |
| 92 | + hidden, |
| 93 | + grad_out_stride, |
| 94 | + rows_stride, |
| 95 | + scores_stride, |
| 96 | + grad_rows_stride, |
| 97 | + grad_scores_stride, |
| 98 | + TOP_K: tl.constexpr, |
| 99 | + K_PADDED: tl.constexpr, |
| 100 | + BLOCK_H: tl.constexpr, |
| 101 | + ): |
| 102 | + token = tl.program_id(0) |
| 103 | + |
| 104 | + slots = tl.arange(0, K_PADDED) |
| 105 | + slot_mask = slots < TOP_K |
| 106 | + |
| 107 | + source_rows = tl.load(inverse_ptr + token * TOP_K + slots, mask=slot_mask, other=-1).to(tl.int64) |
| 108 | + row_valid = slot_mask & (source_rows >= 0) |
| 109 | + safe_rows = tl.where(row_valid, source_rows, 0) |
| 110 | + scores = tl.load(scores_ptr + token * scores_stride + slots, mask=slot_mask, other=0.0).to(tl.float32) |
| 111 | + |
| 112 | + grad_rows_dtype = grad_rows_ptr.dtype.element_ty |
| 113 | + # Keeping one token per program avoids a second reduction pass for scores. |
| 114 | + score_partials = tl.zeros([K_PADDED, BLOCK_H], dtype=tl.float32) |
| 115 | + |
| 116 | + for hidden_start in range(0, hidden, BLOCK_H): |
| 117 | + hidden_offsets = hidden_start + tl.arange(0, BLOCK_H) |
| 118 | + hidden_mask = hidden_offsets < hidden |
| 119 | + block_mask = row_valid[:, None] & hidden_mask[None, :] |
| 120 | + |
| 121 | + upstream = tl.load( |
| 122 | + grad_out_ptr + token * grad_out_stride + hidden_offsets, |
| 123 | + mask=hidden_mask, |
| 124 | + other=0.0, |
| 125 | + ).to(tl.float32) |
| 126 | + |
| 127 | + values = tl.load( |
| 128 | + rows_ptr + safe_rows[:, None] * rows_stride + hidden_offsets[None, :], |
| 129 | + mask=block_mask, |
| 130 | + other=0.0, |
| 131 | + ).to(tl.float32) |
| 132 | + score_partials += values * upstream[None, :] |
| 133 | + |
| 134 | + tl.store( |
| 135 | + grad_rows_ptr + safe_rows[:, None] * grad_rows_stride + hidden_offsets[None, :], |
| 136 | + (upstream[None, :] * scores[:, None]).to(grad_rows_dtype), |
| 137 | + mask=block_mask, |
| 138 | + ) |
| 139 | + |
| 140 | + grad_scores = tl.sum(score_partials, axis=1) |
| 141 | + tl.store( |
| 142 | + grad_scores_ptr + token * grad_scores_stride + slots, |
| 143 | + grad_scores.to(grad_scores_ptr.dtype.element_ty), |
| 144 | + mask=slot_mask, |
| 145 | + ) |
| 146 | + |
| 147 | + |
| 148 | +def is_available() -> bool: |
| 149 | + """Whether this build can run the fused weighted restore at all.""" |
| 150 | + return _TRITON_AVAILABLE and not _IS_ROCM_PYTORCH |
| 151 | + |
| 152 | + |
| 153 | +def assert_supported(rows: torch.Tensor, *, score_apply: str) -> None: |
| 154 | + """Reject unsupported configurations before collectives begin.""" |
| 155 | + if not _TRITON_AVAILABLE: |
| 156 | + raise RuntimeError('combine_impl="fused_weighted_sum" needs Triton, which is not installed in this ' |
| 157 | + "environment. Install Triton, or leave combine_impl unset.") |
| 158 | + if _IS_ROCM_PYTORCH: |
| 159 | + raise RuntimeError('combine_impl="fused_weighted_sum" is not yet supported on ROCm. Leave combine_impl ' |
| 160 | + "unset to run here.") |
| 161 | + if rows.device.type != "cuda": |
| 162 | + raise RuntimeError('combine_impl="fused_weighted_sum" runs CUDA kernels but this layer is on device ' |
| 163 | + f'"{rows.device.type}". Leave combine_impl unset to run here.') |
| 164 | + if rows.dtype not in SUPPORTED_ROW_DTYPES: |
| 165 | + raise RuntimeError('combine_impl="fused_weighted_sum" supports bfloat16, float16, and float32 rows, got ' |
| 166 | + f"{rows.dtype}. Leave combine_impl unset, or use a supported floating-point dtype.") |
| 167 | + if score_apply != "post": |
| 168 | + raise RuntimeError('combine_impl="fused_weighted_sum" folds the routing weight into the top-k reduction, ' |
| 169 | + f'which only exists for score_apply="post", but this layer resolved ' |
| 170 | + f'score_apply="{score_apply}". Leave combine_impl unset.') |
| 171 | + |
| 172 | + |
| 173 | +def _block_hidden(hidden: int, slots: int) -> int: |
| 174 | + """Choose a power-of-two tile within the FP32 register budget.""" |
| 175 | + budget = max(16, _MAX_BLOCK_ELEMENTS // slots) |
| 176 | + return min(_MAX_BLOCK_HIDDEN, budget, max(16, triton.next_power_of_2(hidden))) |
| 177 | + |
| 178 | + |
| 179 | +def _padded_top_k(top_k: int) -> int: |
| 180 | + """Round top-k up to a power of two, which ``tl.arange`` requires.""" |
| 181 | + return max(2, triton.next_power_of_2(top_k)) |
| 182 | + |
| 183 | + |
| 184 | +def _invert_index(index: torch.Tensor, num_inverse_rows: int) -> torch.Tensor: |
| 185 | + """Invert the row permutation produced by sorting routed assignments.""" |
| 186 | + inverse = torch.empty((num_inverse_rows, ), dtype=torch.int32, device=index.device) |
| 187 | + num_indices = index.numel() |
| 188 | + |
| 189 | + grid = (triton.cdiv(num_indices, _INVERT_INDEX_BLOCK), ) |
| 190 | + _invert_index_kernel[grid]( |
| 191 | + index.contiguous(), |
| 192 | + inverse, |
| 193 | + num_indices, |
| 194 | + num_inverse_rows, |
| 195 | + BLOCK=_INVERT_INDEX_BLOCK, |
| 196 | + ) |
| 197 | + return inverse |
| 198 | + |
| 199 | + |
| 200 | +class _FusedWeightedRestore(torch.autograd.Function): |
| 201 | + """Weight rows by their routing score and reduce over top-k in one pass.""" |
| 202 | + |
| 203 | + @staticmethod |
| 204 | + def forward(ctx, combined_rows, top_scores, inverse, top_k): |
| 205 | + combined_rows = combined_rows.contiguous() |
| 206 | + n_tokens, hidden = top_scores.shape[0], combined_rows.shape[-1] |
| 207 | + output = torch.empty((n_tokens, hidden), dtype=combined_rows.dtype, device=combined_rows.device) |
| 208 | + |
| 209 | + ctx.save_for_backward(combined_rows, top_scores, inverse) |
| 210 | + ctx.top_k = top_k |
| 211 | + |
| 212 | + k_padded = _padded_top_k(top_k) |
| 213 | + block_hidden = _block_hidden(hidden, slots=k_padded) |
| 214 | + grid = (n_tokens, triton.cdiv(hidden, block_hidden)) |
| 215 | + _weighted_restore_forward_kernel[grid]( |
| 216 | + combined_rows, |
| 217 | + inverse, |
| 218 | + top_scores, |
| 219 | + output, |
| 220 | + hidden, |
| 221 | + combined_rows.stride(0), |
| 222 | + top_scores.stride(0), |
| 223 | + output.stride(0), |
| 224 | + TOP_K=top_k, |
| 225 | + K_PADDED=k_padded, |
| 226 | + BLOCK_H=block_hidden, |
| 227 | + ) |
| 228 | + return output |
| 229 | + |
| 230 | + @staticmethod |
| 231 | + def backward(ctx, grad_output): |
| 232 | + combined_rows, top_scores, inverse = ctx.saved_tensors |
| 233 | + grad_output = grad_output.contiguous() |
| 234 | + |
| 235 | + grad_rows = torch.empty_like(combined_rows) |
| 236 | + grad_scores = torch.empty_like(top_scores) |
| 237 | + |
| 238 | + n_tokens, hidden = top_scores.shape[0], combined_rows.shape[-1] |
| 239 | + |
| 240 | + k_padded = _padded_top_k(ctx.top_k) |
| 241 | + _weighted_restore_backward_kernel[(n_tokens, )]( |
| 242 | + grad_output, |
| 243 | + combined_rows, |
| 244 | + inverse, |
| 245 | + top_scores, |
| 246 | + grad_rows, |
| 247 | + grad_scores, |
| 248 | + hidden, |
| 249 | + grad_output.stride(0), |
| 250 | + combined_rows.stride(0), |
| 251 | + top_scores.stride(0), |
| 252 | + grad_rows.stride(0), |
| 253 | + grad_scores.stride(0), |
| 254 | + TOP_K=ctx.top_k, |
| 255 | + K_PADDED=k_padded, |
| 256 | + BLOCK_H=_block_hidden(hidden, slots=k_padded), |
| 257 | + ) |
| 258 | + return grad_rows, grad_scores, None, None |
| 259 | + |
| 260 | + |
| 261 | +def fused_weighted_restore( |
| 262 | + combined_rows: torch.Tensor, |
| 263 | + top_scores: torch.Tensor, |
| 264 | + token_indices_sorted: torch.Tensor, |
| 265 | + top_k: int, |
| 266 | + shape: tuple[int, int, int], |
| 267 | +) -> torch.Tensor: |
| 268 | + """Restore ``[T * K, H]`` rows directly to weighted ``[B, S, H]`` output.""" |
| 269 | + bsz, seqlen, hidden = shape |
| 270 | + n_tokens = bsz * seqlen |
| 271 | + expected_rows = n_tokens * top_k |
| 272 | + inverse = _invert_index(token_indices_sorted, expected_rows) |
| 273 | + output = _FusedWeightedRestore.apply(combined_rows, top_scores.contiguous(), inverse, top_k) |
| 274 | + return output.reshape(bsz, seqlen, hidden) |
0 commit comments