Diff Coverage

                 context: Mapping[str, Any] | None = None) -> None:
        """Validate both fused kernels during construction, with module diagnostics."""
        self.clamp_limit = float(module.limit) if isinstance(module, DeepseekV4Experts) else None
        if self.clamp_limit is not None and (not math.isfinite(self.clamp_limit) or self.clamp_limit <= 0):
            raise ValueError(f"{module_fqn or '<root>'}: clipped MXFP8 fusion requires a positive finite limit")
        super().__init__(module=module, module_fqn=module_fqn, context=context)
        try:
            validate_npu_swiglu_runtime(clamp_limit=self.clamp_limit)
        except RuntimeError as exc:

def _fused_quant(projected, quantizer, counts, gradient=None, clamp_limit=None):
    """Keep prequantized operands internal to this Function's GMM calls."""
    options = {
        "quant_dtype": quantizer.quant_dtype,
        "group_index": counts.cumsum(0).to(torch.int64),
        "clamp_limit": clamp_limit,
    }

            clamp_limit: Optional positive clipping limit for DeepSeek experts.
        """
        ctx.shapes = (inputs.shape, gate_up.shape, down.shape)
        ctx.dtype, ctx.device, ctx.quantizer = inputs.dtype, inputs.device, quantizer
        ctx.clamp_limit = clamp_limit
        if inputs.shape[0] == 0:
            ctx.save_for_backward()
            return inputs.new_empty((0, down.shape[-1]))
        input_q = quantizer.quantize(inputs, group_list=counts, group_list_type=1, rowwise=True, colwise=True)

        input_q = quantizer.quantize(inputs, group_list=counts, group_list_type=1, rowwise=True, colwise=True)
        gate_up_q = quantizer.quantize(gate_up.contiguous(), rowwise=True, colwise=True)
        down_q = quantizer.quantize(down.contiguous(), rowwise=True, colwise=True)
        projected = _gmm(input_q, gate_up_q, counts, inputs.dtype)
        activated_q = _fused_quant(projected, quantizer, counts, clamp_limit=clamp_limit)
        output = _gmm(activated_q, down_q, counts, inputs.dtype)
        _save_operands(ctx, counts, projected, (input_q, gate_up_q, activated_q, down_q))
        return output

        needs_input, needs_gate, needs_down = ctx.needs_input_grad[:3]
        if ctx.shapes[0][0] == 0:
            grads = [torch.zeros(shape, device=ctx.device, dtype=ctx.dtype) if needed else None
                     for shape, needed in zip(ctx.shapes, (needs_input, needs_gate, needs_down))]
            return *grads, None, None, None
        counts, projected, (input_q, gate_up_q, activated_q, down_q) = _restore_operands(ctx)
        grad_q = ctx.quantizer.quantize(
            gradient, group_list=counts if needs_down else None, group_list_type=1,
            rowwise=needs_input or needs_gate, colwise=needs_down)

        grad_down = _gmm(activated_q, grad_q, counts, ctx.dtype, "TN") if needs_down else None
        grad_input = grad_gate = None
        if needs_input or needs_gate:
            grad_activated = _gmm(grad_q, down_q, counts, ctx.dtype, "NT")
            grad_projected_q = _fused_quant(projected, ctx.quantizer, counts, grad_activated, ctx.clamp_limit)
            if needs_input:
                grad_input = _gmm(grad_projected_q, gate_up_q, counts, ctx.dtype, "NT")
            if needs_gate:
                grad_gate = _gmm(input_q, grad_projected_q, counts, ctx.dtype, "TN")
        return grad_input, grad_gate, grad_down, None, None, None


def fused_grouped_swiglu(
    inputs: torch.Tensor, gate_up: torch.Tensor, down: torch.Tensor,

            or gate_up.shape[2] != 2 * down.shape[1] or down.shape[2] != inputs.shape[1]):
        raise ValueError("Fused MXFP8 SwiGLU weight dimensions do not match inputs")
    if counts.ndim != 1 or counts.numel() != gate_up.shape[0] or counts.dtype not in (torch.int32, torch.int64):
        raise ValueError("Fused MXFP8 SwiGLU requires one integer count per expert")
    return _FusedGroupedSwiGLU.apply(inputs, gate_up, down, counts, quantizer, clamp_limit)

def _get_clamped_swiglu_ops() -> Any:
    """Load optional CANN registration only when clipped fusion is requested."""
    try:
        import cann_ops_nn  # pylint: disable=C0415,unused-import
    except ImportError as exc:
        raise LowPrecisionCapabilityError(
            "Clipped MXFP8 SwiGLU requires cann_ops_nn (CANN 9.2.1 on Ascend 950PR/950DT)."
        ) from exc
    operators = torch.ops.cann_ops_nn
    names = ("swiglu_group_quant_with_dual_axis", "swiglu_backward_group_quant_with_dual_axis")

    names = ("swiglu_group_quant_with_dual_axis", "swiglu_backward_group_quant_with_dual_axis")
    missing = [name for name in names if not hasattr(operators, name)]
    if missing:
        raise LowPrecisionCapabilityError(f"Clipped MXFP8 SwiGLU requires cann_ops_nn operators {missing}")
    return operators


class MXFP8NpuOps:
    """Adapt the A5 MXFP8 torch_npu operator contract."""

def validate_npu_swiglu_runtime(*, clamp_limit: Optional[float] = None) -> None:
    """Validate fusion without storing runtime modules on a model's quantizer."""
    _get_npu_ops().validate_swiglu_fusion(clamp_limit=clamp_limit)


def validate_npu_gmm_runtime() -> None:
    """Fail during expert setup when the MXFP8 GMM runtime is unavailable."""