« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Fused SwiGLU/quantization within one grouped-expert autograd boundary."""
17from typing import TYPE_CHECKING, Optional
19import torch # pylint: disable=forbidden-backend-import
20from torch.autograd.function import once_differentiable # pylint: disable=forbidden-backend-import
22from hyper_parallel.components.quantization.ops.npu_mxfp8 import mxfp8_grouped_matmul
23from hyper_parallel.components.quantization.tensor import MXFP8Tensor
25if TYPE_CHECKING:
26 from hyper_parallel.components.quantization.quantizers.mxfp8 import MXFP8Quantizer
29def _save_operands(ctx, counts, projected, operands):
30 """Expose physical tensors to saved-tensor hooks, never carrier wrappers."""
31 ctx.operand_metadata = tuple((tuple(q.shape), q.dtype, q.quantizer) for q in operands)
32 buffers = tuple(buffer for q in operands for buffer in (q.row_data, q.row_scale, q.col_data, q.col_scale))
33 ctx.save_for_backward(counts, projected, *buffers)
36def _restore_operands(ctx):
37 """Use the same physical-buffer reconstruction as the upstream Dense/GMM path."""
38 counts, projected, *buffers = ctx.saved_tensors
39 operands = []
40 for index, (shape, dtype, quantizer) in enumerate(ctx.operand_metadata):
41 row, row_scale, col, col_scale = buffers[4 * index:4 * index + 4]
42 operands.append(MXFP8Tensor(shape, dtype, quantizer=quantizer, row_data=row,
43 row_scale=row_scale, col_data=col, col_scale=col_scale))
44 return counts, projected, operands
47def _fused_quant(projected, quantizer, counts, gradient=None):
48 """Keep prequantized operands internal to this Function's GMM calls."""
49 options = {"quant_dtype": quantizer.quant_dtype, "group_index": counts.cumsum(0).to(torch.int64)}
50 if gradient is None:
51 buffers = quantizer.npu_ops.swiglu_mx_quant_dual_axis(projected, **options)
52 shape = (projected.shape[0], projected.shape[1] // 2)
53 else:
54 buffers = quantizer.npu_ops.swiglu_backward_mx_quant_dual_axis(
55 projected, gradient.contiguous(), **options)
56 shape = projected.shape
57 row, row_scale, col, col_scale = buffers
58 return MXFP8Tensor(shape, projected.dtype, quantizer=quantizer,
59 row_data=row, row_scale=row_scale, col_data=col, col_scale=col_scale)
62def _gmm(left, right, counts, dtype, layout="NN"):
63 """Use the existing MXFP8 adapter and count-based expert partitioning."""
64 return mxfp8_grouped_matmul(left, right, layout=layout, group_list=counts,
65 group_type=2 if layout == "TN" else 0,
66 group_list_type=1, output_dtype=dtype)
69class _FusedGroupedSwiGLU(torch.autograd.Function):
70 """Own both GMMs and fusion so external autograd edges remain ordinary tensors."""
72 @staticmethod
73 def forward(
74 ctx: torch.autograd.function.FunctionCtx, inputs: torch.Tensor,
75 gate_up: torch.Tensor, down: torch.Tensor, counts: torch.Tensor, quantizer: "MXFP8Quantizer",
76 ) -> torch.Tensor:
77 """Compute two MXFP8 GMMs with a fused activation/quantization boundary.
79 Args:
80 ctx: Autograd save context.
81 inputs: Expert-major token matrix.
82 gate_up: Canonical gate/up weights.
83 down: Canonical down weights.
84 counts: Per-expert token counts.
85 quantizer: MXFP8 quantizer.
86 """
87 ctx.shapes = (inputs.shape, gate_up.shape, down.shape)
88 ctx.dtype, ctx.device, ctx.quantizer = inputs.dtype, inputs.device, quantizer
89 if inputs.shape[0] == 0:
90 ctx.save_for_backward()
91 return inputs.new_empty((0, down.shape[-1]))
92 input_q = quantizer.quantize(inputs, group_list=counts, group_list_type=1, rowwise=True, colwise=True)
93 gate_up_q = quantizer.quantize(gate_up.contiguous(), rowwise=True, colwise=True)
94 down_q = quantizer.quantize(down.contiguous(), rowwise=True, colwise=True)
95 projected = _gmm(input_q, gate_up_q, counts, inputs.dtype)
96 activated_q = _fused_quant(projected, quantizer, counts)
97 output = _gmm(activated_q, down_q, counts, inputs.dtype)
98 _save_operands(ctx, counts, projected, (input_q, gate_up_q, activated_q, down_q))
99 return output
101 @staticmethod
102 @once_differentiable
103 def backward(
104 ctx: torch.autograd.function.FunctionCtx, gradient: torch.Tensor,
105 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], None, None]:
106 """Fuse SwiGLU derivative/quantization and compute requested first gradients.
108 Args:
109 ctx: Saved forward context.
110 gradient: Output gradient.
111 """
112 needs_input, needs_gate, needs_down = ctx.needs_input_grad[:3]
113 if ctx.shapes[0][0] == 0:
114 grads = [torch.zeros(shape, device=ctx.device, dtype=ctx.dtype) if needed else None
115 for shape, needed in zip(ctx.shapes, (needs_input, needs_gate, needs_down))]
116 return *grads, None, None
117 counts, projected, (input_q, gate_up_q, activated_q, down_q) = _restore_operands(ctx)
118 grad_q = ctx.quantizer.quantize(
119 gradient, group_list=counts if needs_down else None, group_list_type=1,
120 rowwise=needs_input or needs_gate, colwise=needs_down)
121 grad_down = _gmm(activated_q, grad_q, counts, ctx.dtype, "TN") if needs_down else None
122 grad_input = grad_gate = None
123 if needs_input or needs_gate:
124 grad_activated = _gmm(grad_q, down_q, counts, ctx.dtype, "NT")
125 grad_projected_q = _fused_quant(projected, ctx.quantizer, counts, grad_activated)
126 if needs_input:
127 grad_input = _gmm(grad_projected_q, gate_up_q, counts, ctx.dtype, "NT")
128 if needs_gate:
129 grad_gate = _gmm(input_q, grad_projected_q, counts, ctx.dtype, "TN")
130 return grad_input, grad_gate, grad_down, None, None
133def fused_grouped_swiglu(
134 inputs: torch.Tensor, gate_up: torch.Tensor, down: torch.Tensor,
135 counts: torch.Tensor, quantizer: "MXFP8Quantizer",
136) -> torch.Tensor:
137 """Compute expert-major outputs with canonical [E,H,2I]/[E,I,H] weights.
139 Args:
140 inputs: BF16 expert-major token matrix.
141 gate_up: Canonical packed gate/up weights.
142 down: Canonical down weights.
143 counts: Per-local-expert token counts, including zero-token experts.
144 quantizer: MXFP8 recipe with both fused NPU kernels available.
146 Returns:
147 Ordinary BF16 outputs before outer routing weights are applied.
148 """
149 if inputs.dtype not in (torch.bfloat16, torch.float16):
150 raise ValueError("Fused MXFP8 SwiGLU requires BF16 or FP16 inputs")
151 if gate_up.dtype != inputs.dtype or down.dtype != inputs.dtype:
152 raise ValueError("Fused MXFP8 SwiGLU requires matching input and weight dtypes")
153 if inputs.ndim != 2 or gate_up.ndim != 3 or down.ndim != 3:
154 raise ValueError("Fused MXFP8 SwiGLU requires 2D inputs and 3D expert weights")
155 if (gate_up.shape[0] != down.shape[0] or gate_up.shape[1] != inputs.shape[1]
156 or gate_up.shape[2] != 2 * down.shape[1] or down.shape[2] != inputs.shape[1]):
157 raise ValueError("Fused MXFP8 SwiGLU weight dimensions do not match inputs")
158 if counts.ndim != 1 or counts.numel() != gate_up.shape[0] or counts.dtype not in (torch.int32, torch.int64):
159 raise ValueError("Fused MXFP8 SwiGLU requires one integer count per expert")
160 return _FusedGroupedSwiGLU.apply(inputs, gate_up, down, counts, quantizer)