Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / function.py: 92%
214 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +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"""Differentiable autograd function for the low-level MegaMoe ABI."""
17from __future__ import annotations
19from dataclasses import dataclass, replace
20from typing import Any
22import torch
24from hyper_multicore._loader import load_torch_npu
25from hyper_multicore.torch import ops as multicore_ops
27from .plan import MegaMoePlan, record_plan_stream
28from .route import PreparedTopKRoute, RouteMetadata
29from .workspace import MegaMoeWorkspace
32@dataclass(frozen=True)
33class _ForwardIntermediates:
34 """Forward tensors owned by one autograd invocation."""
36 up_proj: Any
37 activation: Any
38 down_proj: Any
41@dataclass(frozen=True)
42class _SavedBackwardState:
43 """Forward tensors and route metadata restored by backward."""
45 dispatch: Any
46 up_proj: Any
47 activation: Any
48 weight1: Any
49 weight2: Any
50 group_list: Any
51 dispatch_src_off: Any
52 dispatch_target_off: Any
53 dispatch_size: Any
54 combine_src_off: Any
55 combine_target_off: Any
56 combine_size: Any
59@dataclass(frozen=True)
60class _BackwardIntermediates:
61 """Backward tensors owned by one autograd invocation."""
63 grad_weight1: Any
64 grad_weight2: Any
65 act_grad: Any
66 swiglu_grad: Any
67 gate_dx: Any
70@dataclass(frozen=True)
71class _ForwardExecution:
72 """Borrowed workspace and event state for one forward launch."""
74 dispatch: Any
75 source: Any
76 combine: Any
77 event_counters: Any
78 gmm_workspace: Any
79 intermediates: _ForwardIntermediates
82@dataclass(frozen=True)
83class _BackwardExecution:
84 """Borrowed workspace and event state for one backward launch."""
86 dispatch: Any
87 grad_x: Any
88 event_counters: Any
89 gmm_workspace: Any
90 swiglu_workspace: Any
91 intermediates: _BackwardIntermediates
94def _workspace_tensor(tensor: Any | None, name: str) -> Any:
95 """Return an initialized workspace tensor or raise an invariant error."""
96 if tensor is None:
97 raise RuntimeError(f"MegaMoe workspace {name} is not initialized.")
98 return tensor
101def _dispatch_and_source(spec: Any, workspace: Any, rows: Any, capacity: int) -> tuple[Any, Any]:
102 """Select transport storage while retaining a source-ready stream dependency."""
103 if getattr(spec, "dispatch_mode", "push") == "push":
104 return _workspace_tensor(workspace.expert_buffer, "expert_buffer"), rows
105 source = _workspace_tensor(workspace.source_buffer, "source_buffer")[:rows.shape[0]]
106 if source.data_ptr() != rows.data_ptr():
107 source.copy_(rows)
108 dispatch = torch.empty((capacity, spec.hidden_size), dtype=rows.dtype, device=rows.device)
109 return dispatch, source
112def _stage_backward_source(ctx: Any, grad_output: Any, permutation_inputs: tuple[Any, ...]) -> tuple[Any, Any]:
113 """Write pull dY into SHMEM before expert scratch allocation."""
114 if getattr(ctx.plan.spec, "dispatch_mode", "push") == "push":
115 return grad_output, None
116 source = _workspace_tensor(ctx.workspace.source_buffer, "source_buffer")[:ctx.source_rows]
117 if not ctx.has_unpermute:
118 source.copy_(grad_output)
119 return source, None
120 mapping, expert_output, probs = permutation_inputs
121 grad_probs = torch.empty_like(probs)
122 if ctx.source_rows:
123 multicore_ops.mega_moe_unpermute_grad_out(expert_output, grad_output, mapping, probs, source, grad_probs)
124 return source, grad_probs.to(ctx.topk_dtype) if ctx.needs_input_grad[7] else None
127def _allocate_forward_intermediates(
128 spec: Any,
129 capacity: int,
130 routed_tokens: Any,
131 dispatch: Any,
132) -> _ForwardIntermediates:
133 """Allocate outputs whose received rows are overwritten before use."""
134 # GMM and SwiGLU cover the received prefix; an empty rank reads no rows.
135 up_proj = torch.empty(
136 (capacity, spec.intermediate_size * 2),
137 dtype=routed_tokens.dtype,
138 device=routed_tokens.device,
139 )
140 activation = torch.empty(
141 (capacity, spec.intermediate_size),
142 dtype=routed_tokens.dtype,
143 device=routed_tokens.device,
144 )
145 return _ForwardIntermediates(up_proj, activation, torch.empty_like(dispatch[:capacity]))
148def _allocate_backward_intermediates(
149 spec: Any,
150 capacity: int,
151 grad_output: Any,
152 weight1: Any,
153 weight2: Any,
154 reusable_dispatch: Any | None = None,
155) -> _BackwardIntermediates:
156 """Allocate overwritten activations and zero-safe expert gradients."""
157 grad_weight1 = torch.zeros_like(weight1)
158 grad_weight2 = torch.zeros_like(weight2)
159 act_grad = torch.empty(
160 (capacity, spec.intermediate_size),
161 dtype=grad_output.dtype,
162 device=grad_output.device,
163 )
164 swiglu_grad = torch.empty(
165 (capacity, spec.intermediate_size * 2),
166 dtype=grad_output.dtype,
167 device=grad_output.device,
168 )
169 gate_dx = reusable_dispatch if reusable_dispatch is not None else torch.empty(
170 (capacity, spec.hidden_size),
171 dtype=grad_output.dtype,
172 device=grad_output.device,
173 )
174 # Empty experts must retain exact zero gradients on the baseline kernel.
175 return _BackwardIntermediates(grad_weight1, grad_weight2, act_grad, swiglu_grad, gate_dx)
178def _native_weight_operand(weight: torch.Tensor, plan: MegaMoePlan) -> torch.Tensor:
179 """Prepare one expert operand without copying native EOI storage."""
180 if plan.spec.expert_weight_layout == "EOI":
181 return weight
182 return weight.contiguous()
185def _saved_backward_state(saved_tensors: tuple[Any, ...]) -> _SavedBackwardState:
186 """Restore named kernel state from the flat autograd tensor tuple."""
187 return _SavedBackwardState(*saved_tensors[:12])
190def _prepare_forward_execution(
191 workspace: MegaMoeWorkspace,
192 plan: MegaMoePlan,
193 routed_tokens: Any,
194 capacity: int,
195) -> _ForwardExecution:
196 """Resolve workspace, event counters, and intermediate tensors for forward."""
197 dispatch, source = _dispatch_and_source(plan.spec, workspace, routed_tokens, capacity)
198 combine = _workspace_tensor(workspace.routed_buffer, "routed_buffer")
199 return _ForwardExecution(
200 dispatch=dispatch,
201 source=source,
202 combine=combine,
203 event_counters=workspace.prepare_event_counters(forward=True),
204 gmm_workspace=_workspace_tensor(workspace.gmm_workspace, "gmm_workspace"),
205 intermediates=_allocate_forward_intermediates(
206 plan.spec,
207 capacity,
208 routed_tokens,
209 dispatch,
210 ),
211 )
214def _prepare_backward_execution(
215 workspace: MegaMoeWorkspace,
216 plan: MegaMoePlan,
217 saved: _SavedBackwardState,
218 grad_output: Any,
219) -> _BackwardExecution:
220 """Resolve workspace, event counters, and intermediate tensors for backward."""
221 dispatch = (torch.empty_like(saved.dispatch) if plan.spec.dispatch_mode == "pull"
222 else _workspace_tensor(workspace.expert_buffer, "expert_buffer"))
223 grad_x = _workspace_tensor(workspace.routed_buffer, "routed_buffer")
224 return _BackwardExecution(
225 dispatch=dispatch,
226 grad_x=grad_x,
227 event_counters=workspace.prepare_event_counters(forward=False),
228 gmm_workspace=_workspace_tensor(workspace.gmm_workspace, "gmm_workspace"),
229 swiglu_workspace=_workspace_tensor(
230 workspace.swiglu_grad_workspace,
231 "swiglu_grad_workspace",
232 ),
233 intermediates=_allocate_backward_intermediates(
234 plan.spec,
235 saved.dispatch.shape[0],
236 grad_output,
237 saved.weight1,
238 saved.weight2,
239 dispatch[:saved.dispatch.shape[0]] if plan.reuse_backward_dispatch else None,
240 ),
241 )
244def _restore_input_gradient(ctx: Any, grad_x: Any, permutation_inputs: tuple[Any, ...]) -> Any:
245 """Return an owned input gradient while the workspace lease is held."""
246 if not ctx.needs_input_grad[0]:
247 return None
248 grad_x = grad_x[:ctx.source_rows]
249 if not ctx.has_permutation:
250 return grad_x.clone()
251 (unpermute_mapping,) = permutation_inputs
252 spec = ctx.plan.spec
253 if ctx.token_rows == 0:
254 return grad_x[:0].clone()
255 # Consume the shared gradient before release records completion.
256 # The permutation gradient owns its reduced [T, H] output.
257 return multicore_ops.moe_token_permute_grad(
258 grad_x, unpermute_mapping, ctx.token_rows, spec.top_k
259 )
262def _save_forward_state(
263 ctx: Any,
264 plan: MegaMoePlan,
265 workspace: MegaMoeWorkspace,
266 saved_dispatch: Any,
267 up_proj: Any,
268 activation: Any,
269 weight1: Any,
270 weight2: Any,
271 metadata: RouteMetadata,
272 permutation_inputs: tuple[Any, ...] = (),
273) -> None:
274 """Save owned forward tensors and route state for backward."""
275 ctx.plan = plan
276 ctx.workspace = workspace
277 ctx.save_for_backward(
278 saved_dispatch,
279 up_proj,
280 activation,
281 weight1,
282 weight2,
283 metadata.group_list,
284 metadata.dispatch_src_off,
285 metadata.dispatch_target_off,
286 metadata.dispatch_size,
287 metadata.combine_src_off,
288 metadata.combine_target_off,
289 metadata.combine_size,
290 *permutation_inputs,
291 )
292 workspace.track_graph(ctx)
295def _release_completed_graph(ctx: Any, workspace: MegaMoeWorkspace) -> None:
296 """Keep retained graphs alive; older engines conservatively wait for context GC."""
297 # The engine's retention flag avoids unpacking saved tensors or guessing from output references.
298 keep_graph = getattr(torch._C._autograd, "_get_current_graph_task_keep_graph", None) # pylint: disable=protected-access
299 if keep_graph is not None and not keep_graph():
300 workspace.release_graph(ctx)
303def _launch_forward_kernel(
304 plan: MegaMoePlan,
305 metadata: RouteMetadata,
306 routed_tokens: Any,
307 weight1: Any,
308 weight2: Any,
309 execution: _ForwardExecution,
310) -> None:
311 """Launch the internal forward ABI with prepared buffers and metadata."""
312 spec = plan.spec
313 kernel_parameters = (
314 execution.dispatch,
315 metadata.dispatch_target_off * spec.hidden_size,
316 _native_source(routed_tokens),
317 metadata.dispatch_src_off * spec.hidden_size,
318 metadata.dispatch_size * spec.hidden_size,
319 weight1,
320 metadata.group_list,
321 execution.intermediates.up_proj,
322 execution.intermediates.activation,
323 weight2,
324 metadata.group_list,
325 execution.intermediates.down_proj,
326 execution.combine,
327 metadata.combine_target_off * spec.hidden_size,
328 metadata.combine_src_off * spec.hidden_size,
329 metadata.combine_size * spec.hidden_size,
330 )
331 if not plan.fwd_diagnostics.is_active:
332 record_plan_stream(plan, plan.fwd_runtime_config)
333 multicore_ops.mega_moe(
334 *kernel_parameters,
335 execution.gmm_workspace,
336 plan.up_proj_tiling,
337 plan.swiglu_tiling,
338 plan.down_proj_tiling,
339 plan.fwd_runtime_config,
340 execution.event_counters,
341 spec.rank_id,
342 spec.ep_size,
343 spec.num_experts,
344 spec.hidden_size,
345 getattr(spec, "plan_tokens", spec.local_num_tokens),
346 )
347 return
349 with plan.fwd_diagnostics.launch(
350 kernel_parameters,
351 direction="forward",
352 fallback_event_counters=execution.event_counters,
353 ) as diagnostic_call:
354 record_plan_stream(plan, diagnostic_call.runtime_config)
355 multicore_ops.mega_moe_with_diagnostic_buffer(
356 *kernel_parameters,
357 execution.gmm_workspace,
358 plan.up_proj_tiling,
359 plan.swiglu_tiling,
360 plan.down_proj_tiling,
361 diagnostic_call.runtime_config,
362 diagnostic_call.event_counters,
363 diagnostic_call.diagnostic_buffer,
364 spec.rank_id,
365 spec.ep_size,
366 spec.num_experts,
367 spec.hidden_size,
368 getattr(spec, "plan_tokens", spec.local_num_tokens),
369 )
372def _launch_backward_kernel(
373 plan: MegaMoePlan,
374 saved: _SavedBackwardState,
375 grad_output: Any,
376 execution: _BackwardExecution,
377) -> None:
378 """Launch the internal backward ABI with prepared buffers and metadata."""
379 spec = plan.spec
380 gradients = execution.intermediates
381 kernel_parameters = (
382 execution.dispatch,
383 saved.dispatch_target_off * spec.hidden_size,
384 _native_source(grad_output),
385 saved.dispatch_src_off * spec.hidden_size,
386 saved.dispatch_size * spec.hidden_size,
387 saved.activation,
388 gradients.grad_weight2,
389 saved.weight2,
390 gradients.act_grad,
391 saved.up_proj,
392 gradients.swiglu_grad,
393 saved.weight1,
394 gradients.gate_dx,
395 execution.grad_x,
396 saved.combine_target_off * spec.hidden_size,
397 saved.combine_src_off * spec.hidden_size,
398 saved.combine_size * spec.hidden_size,
399 saved.dispatch,
400 gradients.grad_weight1,
401 saved.group_list,
402 )
403 if not plan.bwd_diagnostics.is_active:
404 record_plan_stream(plan, plan.bwd_runtime_config)
405 multicore_ops.mega_moe_grad(
406 *kernel_parameters,
407 plan.act_grad_tiling,
408 plan.gate_grad_tiling,
409 plan.w1_grad_tiling,
410 plan.w2_grad_tiling,
411 plan.swiglu_grad_tiling,
412 execution.gmm_workspace,
413 execution.swiglu_workspace,
414 plan.bwd_runtime_config,
415 execution.event_counters,
416 spec.rank_id,
417 spec.ep_size,
418 spec.num_experts,
419 spec.hidden_size,
420 getattr(spec, "plan_tokens", spec.local_num_tokens),
421 )
422 return
424 with plan.bwd_diagnostics.launch(
425 kernel_parameters,
426 direction="backward",
427 fallback_event_counters=execution.event_counters,
428 ) as diagnostic_call:
429 record_plan_stream(plan, diagnostic_call.runtime_config)
430 multicore_ops.mega_moe_grad_with_diagnostic_buffer(
431 *kernel_parameters,
432 plan.act_grad_tiling,
433 plan.gate_grad_tiling,
434 plan.w1_grad_tiling,
435 plan.w2_grad_tiling,
436 plan.swiglu_grad_tiling,
437 execution.gmm_workspace,
438 execution.swiglu_workspace,
439 diagnostic_call.runtime_config,
440 diagnostic_call.event_counters,
441 diagnostic_call.diagnostic_buffer,
442 spec.rank_id,
443 spec.ep_size,
444 spec.num_experts,
445 spec.hidden_size,
446 getattr(spec, "plan_tokens", spec.local_num_tokens),
447 )
450def _native_source(tensor: torch.Tensor) -> torch.Tensor:
451 """Keep a non-null ABI pointer and preserve empty symmetric source storage."""
452 if tensor.shape[0] == 0:
453 row_elements = tensor.shape[1]
454 required_bytes = (tensor.storage_offset() + row_elements) * tensor.element_size()
455 if tensor.untyped_storage().nbytes() >= required_bytes:
456 # Pull derives remote addresses from this local symmetric base even on empty ranks.
457 return tensor.as_strided((1, row_elements), (row_elements, 1))
458 return tensor.new_empty((1, tensor.shape[1]))
459 return tensor.contiguous()
462# Torch declares variadic autograd hooks; concrete functions use operator-specific signatures.
463class _MegaMoeFunction(torch.autograd.Function): # pylint: disable=abstract-method,arguments-differ
464 """Autograd bridge for the in-place forward/backward custom ops."""
466 @staticmethod
467 # pylint: disable-next=arguments-differ
468 def forward(
469 ctx: Any,
470 routed_tokens: Any,
471 weight1: Any,
472 weight2: Any,
473 route: RouteMetadata,
474 plan: MegaMoePlan,
475 workspace: MegaMoeWorkspace,
476 permutation: tuple[Any, Any, Any] | None,
477 topk_weights: torch.Tensor | None,
478 workspace_claimed: bool,
479 ) -> Any:
480 """Launch the legacy forward op and save owned backward inputs.
482 Args:
483 ctx: Autograd context for saved backward state.
484 routed_tokens: Expert-major token rows.
485 weight1: Local gate and up-projection weights.
486 weight2: Local down-projection weights.
487 route: Route offsets, sizes and group metadata.
488 plan: Shape-specific forward and backward descriptors.
489 workspace: Reusable directional execution buffers.
490 permutation: Optional routed rows, expert IDs and inverse mapping
491 when the first tensor argument contains original token rows.
492 topk_weights: Optional Router weights for pull source-output gradients.
493 workspace_claimed: Whether the caller owns the forward workspace lease.
495 Returns:
496 Owned expert-major output rows.
497 """
498 spec = plan.spec
499 if topk_weights is not None and (spec.dispatch_mode != "pull" or permutation is None):
500 raise ValueError("Integrated output unpermutation requires pull dispatch and an input permutation.")
501 metadata = route
502 permutation_inputs = ()
503 if permutation is not None:
504 routed_tokens, _, unpermute_mapping = permutation
505 if ctx.needs_input_grad[0] or topk_weights is not None:
506 permutation_inputs = (unpermute_mapping,)
507 ctx.source_rows = routed_tokens.shape[0]
508 ctx.token_rows = ctx.source_rows // spec.top_k
509 ctx.has_permutation = permutation is not None
510 ctx.has_unpermute = topk_weights is not None
511 if workspace_claimed:
512 if not workspace.in_use:
513 raise RuntimeError("MegaMoe forward requires the caller-held workspace lease.")
514 else:
515 if not spec.auto_tokens:
516 workspace.ensure(spec, routed_tokens.dtype, routed_tokens.device)
517 workspace.claim()
518 execution = None
519 try:
520 # Dispatch and combine overwrite disjoint route ranges before consumers run.
521 capacity = metadata.expert_capacity
522 execution = _prepare_forward_execution(
523 workspace,
524 plan,
525 routed_tokens,
526 capacity,
527 )
528 _launch_forward_kernel(
529 plan,
530 metadata,
531 execution.source,
532 weight1,
533 weight2,
534 execution,
535 )
536 output = execution.combine[:ctx.source_rows].clone()
537 if spec.dispatch_mode == "pull":
538 saved_dispatch = execution.dispatch
539 else:
540 # Combine has consumed down_proj before SHMEM dispatch is reused.
541 execution.intermediates.down_proj.copy_(execution.dispatch[:capacity])
542 saved_dispatch = execution.intermediates.down_proj
543 if ctx.has_unpermute:
544 probs = topk_weights.float().contiguous()
545 ctx.topk_dtype = topk_weights.dtype
546 permutation_inputs += (output, probs)
547 output = (load_torch_npu().npu_moe_token_unpermute(output, unpermute_mapping, probs=probs)
548 if ctx.source_rows else output[:0])
549 _save_forward_state(
550 ctx,
551 plan,
552 workspace,
553 saved_dispatch,
554 execution.intermediates.up_proj,
555 execution.intermediates.activation,
556 weight1,
557 weight2,
558 metadata,
559 permutation_inputs,
560 )
561 return output
562 finally:
563 if not workspace_claimed:
564 workspace.release()
566 @staticmethod
567 # pylint: disable-next=arguments-differ
568 def backward(ctx: Any, grad_output: Any) -> tuple[Any, ...]:
569 """Launch the legacy backward op into zero-safe gradient buffers.
571 Args:
572 ctx: Autograd context populated by ``forward``.
573 grad_output: Expert-major output gradient rows.
575 Returns:
576 Gradients aligned with the differentiable forward inputs.
577 """
578 plan = ctx.plan
579 workspace = ctx.workspace
580 # Non-reentrant checkpoint hooks allow each saved tensor to be unpacked only once.
581 saved_tensors = ctx.saved_tensors
582 ctx.maybe_clear_saved_tensors()
583 saved = _saved_backward_state(saved_tensors)
584 permutation_inputs = saved_tensors[12:]
585 del saved_tensors
586 workspace.claim()
587 try:
588 source, grad_topk_weights = _stage_backward_source(ctx, grad_output, permutation_inputs)
589 permutation_inputs = permutation_inputs[:1]
590 execution = _prepare_backward_execution(
591 workspace,
592 plan,
593 saved,
594 grad_output,
595 )
596 _launch_backward_kernel(
597 plan,
598 saved,
599 source,
600 execution,
601 )
602 # Release ordinary scratch before allocating the owned token gradient;
603 # SHMEM remains leased until the current-stream work is recorded.
604 execution = replace(execution, dispatch=None, intermediates=_BackwardIntermediates(
605 execution.intermediates.grad_weight1,
606 execution.intermediates.grad_weight2,
607 None, None, None,
608 ))
609 grad_input = _restore_input_gradient(ctx, execution.grad_x, permutation_inputs)
610 _release_completed_graph(ctx, workspace)
611 return (
612 grad_input,
613 execution.intermediates.grad_weight1,
614 execution.intermediates.grad_weight2,
615 None,
616 None,
617 None,
618 None,
619 grad_topk_weights,
620 None,
621 )
622 finally:
623 workspace.release()
626def execute_mega_moe(
627 routed_tokens: Any,
628 weight1: Any,
629 weight2: Any,
630 route: RouteMetadata,
631 plan: MegaMoePlan,
632 workspace: MegaMoeWorkspace,
633) -> Any:
634 """Execute expert-major rows through the differentiable legacy kernels.
636 Args:
637 routed_tokens: Expert-major token rows.
638 weight1: Local gate and up-projection weights in the plan-selected layout.
639 weight2: Local down-projection weights in the plan-selected layout.
640 route: Route offsets, sizes and group metadata.
641 plan: Shape-specific forward and backward descriptors.
642 workspace: Reusable directional execution buffers.
644 Returns:
645 Expert-major output rows with independent storage.
646 """
647 # Historical input-output views are materialized outside Function.forward so
648 # gradients reach the caller's original weights. Output-input storage is native.
649 return _MegaMoeFunction.apply(
650 routed_tokens,
651 _native_weight_operand(weight1, plan),
652 _native_weight_operand(weight2, plan),
653 route,
654 plan,
655 workspace,
656 None,
657 None,
658 False,
659 )
662def execute_mega_moe_with_permutation(
663 hidden_states: torch.Tensor,
664 topk_ids: torch.Tensor,
665 weight1: torch.Tensor,
666 weight2: torch.Tensor,
667 route: PreparedTopKRoute,
668 plan: MegaMoePlan,
669 workspace: MegaMoeWorkspace,
670 *,
671 topk_weights: torch.Tensor | None = None,
672 workspace_claimed: bool = False,
673) -> torch.Tensor:
674 """Include input permutation backward within the workspace lease.
676 Args:
677 hidden_states: Original flattened token rows.
678 topk_ids: Expert IDs used to prepare the routed rows.
679 weight1: Local gate and up-projection weights in the plan-selected layout.
680 weight2: Local down-projection weights in the plan-selected layout.
681 route: Rows and metadata prepared without recording autograd operations.
682 plan: Shape-specific native descriptors.
683 workspace: Reusable communication buffers.
684 topk_weights: Optional Router probabilities for integrated pull output unpermutation.
685 workspace_claimed: Borrow a caller-held forward lease without releasing it.
686 Backward always acquires its own lease on the original workspace.
688 Returns:
689 Owned expert-major rows, or token-major rows when topk_weights is provided,
690 differentiable with respect to input tokens, expert weights and Router probabilities.
691 """
692 return _MegaMoeFunction.apply(
693 hidden_states,
694 _native_weight_operand(weight1, plan),
695 _native_weight_operand(weight2, plan),
696 route.metadata,
697 plan,
698 workspace,
699 (route.routed_tokens, topk_ids, route.unpermute_mapping),
700 topk_weights,
701 workspace_claimed,
702 )