Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / mxfp8_projection.py: 37%

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"""Non-fused MXFP8 projection with physical-tensor autograd saves.""" 

16 

17from typing import TYPE_CHECKING, Optional 

18 

19import torch 

20from torch.autograd.function import once_differentiable 

21 

22from hyper_parallel.components.quantization.functional._saved_quantized import ( 

23 save_quantized_operands, restore_quantized_operands, 

24) 

25from hyper_parallel.components.quantization.ops.npu_mxfp8 import mxfp8_grouped_matmul 

26 

27if TYPE_CHECKING: 

28 from hyper_parallel.components.quantization.quantizers.mxfp8 import MXFP8Quantizer 

29 

30 

31def _gmm(left, right, counts, dtype, layout="NN"): 

32 """Dispatch the existing MXFP8 grouped operator using per-expert counts.""" 

33 return mxfp8_grouped_matmul(left, right, layout=layout, group_list=counts, 

34 group_type=2 if layout == "TN" else 0, 

35 group_list_type=1, output_dtype=dtype) 

36 

37 

38class _MXFP8Projection(torch.autograd.Function): 

39 """Local projection with physical saves; shared GroupedLinear stays unchanged.""" 

40 

41 @staticmethod 

42 def forward( 

43 ctx: torch.autograd.function.FunctionCtx, inputs: torch.Tensor, weight: torch.Tensor, 

44 counts: torch.Tensor, quantizer: "MXFP8Quantizer", 

45 ) -> torch.Tensor: 

46 """Quantize canonical weights and save only requested backward operands. 

47 

48 Args: 

49 ctx: Autograd save context. 

50 inputs: Expert-major token matrix. 

51 weight: Canonical projection weights. 

52 counts: Per-expert token counts. 

53 quantizer: MXFP8 quantizer. 

54 """ 

55 ctx.shapes = (inputs.shape, weight.shape) 

56 ctx.dtype, ctx.device, ctx.quantizer = inputs.dtype, inputs.device, quantizer 

57 needs_input, needs_weight = ctx.needs_input_grad[:2] 

58 if inputs.shape[0] == 0: 

59 ctx.save_for_backward() 

60 return inputs.new_empty((0, weight.shape[-1])) 

61 xq = quantizer.quantize(inputs, group_list=counts if needs_weight else None, group_list_type=1, 

62 rowwise=True, colwise=needs_weight) 

63 wq = quantizer.quantize(weight.contiguous(), rowwise=True, colwise=True) 

64 output = _gmm(xq, wq, counts, inputs.dtype) 

65 save_quantized_operands(ctx, xq if needs_weight else None, wq if needs_input else None, counts) 

66 return output 

67 

68 @staticmethod 

69 @once_differentiable 

70 def backward( 

71 ctx: torch.autograd.function.FunctionCtx, gradient: torch.Tensor, 

72 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], None, None]: 

73 """Reconstruct temporary operands on every backward, including retained graphs. 

74 

75 Args: 

76 ctx: Saved forward context. 

77 gradient: Output gradient. 

78 """ 

79 needs_input, needs_weight = ctx.needs_input_grad[:2] 

80 if ctx.shapes[0][0] == 0: 

81 grads = [torch.zeros(shape, device=ctx.device, dtype=ctx.dtype) if needed else None 

82 for shape, needed in zip(ctx.shapes, (needs_input, needs_weight))] 

83 return *grads, None, None 

84 counts, xq, wq = restore_quantized_operands(ctx) 

85 gq = ctx.quantizer.quantize(gradient.contiguous(), group_list=counts if needs_weight else None, 

86 group_list_type=1, rowwise=needs_input, colwise=needs_weight) 

87 dx = _gmm(gq, wq, counts, ctx.dtype, "NT") if needs_input else None 

88 dw = _gmm(xq, gq, counts, ctx.dtype, "TN") if needs_weight else None 

89 return dx, dw, None, None 

90 

91 

92def mxfp8_projection(inputs: torch.Tensor, weight: torch.Tensor, counts: torch.Tensor, 

93 quantizer: "MXFP8Quantizer") -> torch.Tensor: 

94 """Apply a local expert-major projection with canonical [E,K,N] weights. 

95 

96 Args: 

97 inputs: Expert-major token matrix. 

98 weight: Canonical expert projection weights. 

99 counts: Token count for each expert. 

100 quantizer: MXFP8 operand quantizer. 

101 

102 Returns: 

103 Unweighted expert-major outputs. 

104 """ 

105 return _MXFP8Projection.apply(inputs, weight, counts, quantizer)