From b3356312a69a948130325f16744ca0874ecf8770 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 14 Oct 2026 14:53:24 +0200 Subject: [PATCH 01/14] [PyTorch] Add torch.compile support for fused MoE routing Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 255 +++++++++++++++++++- transformer_engine/pytorch/router.py | 340 ++++++++++++++++++++++++++- 2 files changed, 593 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index c6c82fa3d78..f87fee70efa 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -6,6 +6,7 @@ import pytest import torch +from torch._dynamo.utils import counters from transformer_engine.pytorch.router import ( QBHistogramMode, @@ -41,7 +42,11 @@ def test_sigmoid_bias_preserves_tiny_scores(topk, histogram_mode, layout, mixed_ bin_bounds = torch.tensor([-1.0, 1.0], device="cuda") histogram = torch.zeros(num_experts, num_bins, dtype=torch.int32, device="cuda") reference_probs, reference_map = topk_score_function_pytorch( - logits, topk, scaling_factor=3.16, score_function="sigmoid", expert_bias=expert_bias + logits, + topk, + scaling_factor=3.16, + score_function="sigmoid", + expert_bias=expert_bias, ) kwargs = {} if histogram_mode is not None: @@ -111,6 +116,254 @@ def _get_tolerances(dtype: torch.dtype, num_experts: int): return atol, rtol +def _compile_logits(dtype): + logits = torch.linspace(-2, 2, 32, device="cuda") + logits = logits + torch.linspace(-0.1, 0.1, 64, device="cuda").unsqueeze(1) + return logits.reshape(2, 32, 32).to(dtype).requires_grad_(True) + + +def _check_compile_reference(actual, expected, dtype): + atol, rtol = { + torch.float32: (1e-6, 2e-5), + torch.float16: (2e-4, 2e-3), + torch.bfloat16: (2e-3, 2e-2), + }[dtype] + torch.testing.assert_close(actual.float(), expected.float(), atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize("routing_map_format", [RoutingMapFormat.BYTEMAP, "bitmap_u8"]) +@pytest.mark.parametrize( + "score_function,use_pre_softmax,grouped", + [ + ("softmax", True, False), + ("softmax", False, True), + ("sigmoid", False, True), + ("sqrtsoftplus", False, False), + ], +) +def test_topk_compile(dtype, mode, routing_map_format, score_function, use_pre_softmax, grouped): + torch._dynamo.reset() + counters.clear() + bias = None if score_function == "softmax" else torch.linspace(0.3, 0, 32, device="cuda") + options = dict( + topk=4, + use_pre_softmax=use_pre_softmax, + num_groups=4 if grouped else None, + group_topk=2 if grouped else None, + scaling_factor=1.2, + score_function=score_function, + expert_bias=bias, + ) + + def router(logits): + return fused_topk_with_score_function( + logits, **options, routing_map_format=routing_map_format + ) + + compiled = torch.compile(router, fullgraph=True, mode=mode) + for reverse in (False, True, False): + torch.compiler.cudagraph_mark_step_begin() + logits = _compile_logits(dtype) + if reverse: + logits = logits.flip(-1).detach().requires_grad_(True) + eager_logits = logits.detach().clone().requires_grad_(True) + ref_logits = logits.detach().float().reshape(-1, 32).requires_grad_(True) + probs, routing_map = compiled(logits) + eager_probs, eager_map = router(eager_logits) + ref_probs, ref_map = topk_score_function_pytorch(ref_logits, **options) + ref_map = ref_map.reshape(logits.shape) + if routing_map_format == "bitmap_u8": + ref_map = _bytemap_to_bitmap_u8(ref_map) + assert not routing_map.requires_grad + torch.testing.assert_close(routing_map, eager_map) + torch.testing.assert_close(routing_map, ref_map) + torch.testing.assert_close(probs, eager_probs, atol=0, rtol=0) + _check_compile_reference(probs, ref_probs.reshape(logits.shape), dtype) + grad = torch.randn((*probs.shape, 2), dtype=dtype, device="cuda")[..., 0] + assert not grad.is_contiguous() + (actual_grad,) = torch.autograd.grad(probs, logits, grad) + (eager_grad,) = torch.autograd.grad(eager_probs, eager_logits, grad) + (ref_grad,) = torch.autograd.grad(ref_probs, ref_logits, grad.float().reshape(-1, 32)) + torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) + _check_compile_reference(actual_grad, ref_grad.reshape(logits.shape), dtype) + with torch.no_grad(): + torch.compiler.cudagraph_mark_step_begin() + actual = compiled(logits) + expected = router(logits) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if mode == "reduce-overhead": + assert not counters["inductor"]["cudagraph_skips"] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize("routing_map_format", [RoutingMapFormat.BYTEMAP, "bitmap_u8"]) +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) +def test_scores_for_aux_loss_compile(dtype, mode, routing_map_format, score_function): + torch._dynamo.reset() + counters.clear() + + def scores(logits): + return fused_compute_score_for_moe_aux_loss(logits, 4, score_function, routing_map_format) + + compiled = torch.compile(scores, fullgraph=True, mode=mode) + for reverse in (False, True, False): + torch.compiler.cudagraph_mark_step_begin() + logits = _compile_logits(dtype) + if reverse: + logits = logits.flip(-1).detach().requires_grad_(True) + eager_logits = logits.detach().clone().requires_grad_(True) + ref_logits = logits.detach().float().reshape(-1, 32).requires_grad_(True) + routing_map, output = compiled(logits) + eager_map, eager_output = scores(eager_logits) + ref_map, ref_output = compute_scores_for_aux_loss_pytorch(ref_logits, 4, score_function) + ref_map = ref_map.reshape(logits.shape) + if routing_map_format == "bitmap_u8": + ref_map = _bytemap_to_bitmap_u8(ref_map) + assert output.dtype == torch.float32 + assert not routing_map.requires_grad + torch.testing.assert_close(routing_map, eager_map) + torch.testing.assert_close(routing_map, ref_map) + torch.testing.assert_close(output, eager_output, atol=0, rtol=0) + _check_compile_reference(output, ref_output.reshape(logits.shape), torch.float32) + grad = torch.randn((*output.shape, 2), device="cuda")[..., 0] + (actual_grad,) = torch.autograd.grad(output, logits, grad) + (eager_grad,) = torch.autograd.grad(eager_output, eager_logits, grad) + (ref_grad,) = torch.autograd.grad(ref_output, ref_logits, grad.reshape(-1, 32)) + assert actual_grad.dtype == dtype + torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) + _check_compile_reference(actual_grad, ref_grad.reshape(logits.shape), dtype) + with torch.no_grad(): + torch.compiler.cudagraph_mark_step_begin() + torch.testing.assert_close(compiled(logits), scores(logits), atol=0, rtol=0) + if mode == "reduce-overhead": + assert not counters["inductor"]["cudagraph_skips"] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize("tensor_count", [False, True]) +def test_aux_loss_compile(dtype, mode, tensor_count): + torch._dynamo.reset() + counters.clear() + counts = torch.arange(1, 33, device="cuda", dtype=torch.int32) + total = torch.tensor(64, device="cuda", dtype=torch.int64) if tensor_count else 64 + + def loss(probs, total_num_tokens): + return fused_moe_aux_loss(probs, counts, total_num_tokens, 16, 4, 0.01) + + compiled = torch.compile(loss, fullgraph=True, mode=mode) + for new_total in (64, 128, 47): + torch.compiler.cudagraph_mark_step_begin() + if tensor_count: + total.fill_(new_total) + else: + total = new_total + probs = _compile_logits(dtype).detach().reshape(-1, 32).requires_grad_(True) + eager_probs = probs.detach().clone().requires_grad_(True) + ref_probs = probs.detach().float().requires_grad_(True) + output = compiled(probs, total) + eager_output = loss(eager_probs, total) + ref_output = aux_loss_pytorch(ref_probs, counts, new_total, 4, 16, 0.01) + assert output.shape == () and output.dtype == dtype + # The native loss reduction accumulates with FP32 atomics. + torch.testing.assert_close(output, eager_output) + _check_compile_reference(output, ref_output, dtype) + grad = output.new_tensor(0.7) + (actual_grad,) = torch.autograd.grad(output, probs, grad) + (eager_grad,) = torch.autograd.grad(eager_output, eager_probs, grad) + (ref_grad,) = torch.autograd.grad(ref_output, ref_probs, grad.float()) + torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) + _check_compile_reference(actual_grad, ref_grad, dtype) + with torch.no_grad(): + torch.compiler.cudagraph_mark_step_begin() + torch.testing.assert_close(compiled(probs, total), loss(probs, total)) + if mode == "reduce-overhead": + assert not counters["inductor"]["cudagraph_skips"] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_router_compile_dynamic(dtype): + torch._dynamo.reset() + counts = torch.arange(1, 131, device="cuda", dtype=torch.int32) + + def router_loss(logits): + probs, routing_map = fused_topk_with_score_function( + logits, 4, False, None, None, None, "sigmoid", None, "bitmap_u8" + ) + _, scores = fused_compute_score_for_moe_aux_loss(logits, 4, "sigmoid", 1) + aux_loss = fused_moe_aux_loss(scores, counts, logits.shape[0], 130, 4, 0.01) + return probs.float().square().sum() + aux_loss, routing_map + + compiled = torch.compile(router_loss, fullgraph=True, dynamic=True) + for num_tokens in (16, 33, 64): + logits = torch.randn(num_tokens, 130, device="cuda", dtype=dtype, requires_grad=True) + eager_logits = logits.detach().clone().requires_grad_(True) + actual, routing_map = compiled(logits) + expected, eager_map = router_loss(eager_logits) + assert routing_map.shape == (num_tokens, 17) + torch.testing.assert_close(routing_map, eager_map) + torch.testing.assert_close(actual, expected) + (actual_grad,) = torch.autograd.grad(actual, logits) + (eager_grad,) = torch.autograd.grad(expected, eager_logits) + _check_compile_reference(actual_grad, eager_grad, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("routing_map_format", [0, 1]) +def test_router_compile_opcheck(dtype, routing_map_format): + from transformer_engine.pytorch import router + + logits = torch.randn(16, 130, device="cuda", dtype=dtype, requires_grad=True) + torch.library.opcheck( + router._topk_forward, + (logits, 4, False, None, None, None, "sigmoid", None, routing_map_format), + ) + torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) + counts = torch.ones(130, device="cuda", dtype=torch.int32) + count = torch.tensor(16, device="cuda", dtype=torch.int64) + for token_count in (None, count): + torch.library.opcheck( + router._aux_loss_forward, (logits, counts, 16, token_count, 130, 4, 0.01) + ) + + +@pytest.mark.parametrize("qb", [False, True]) +def test_router_compile_eager_fallback(qb): + torch._dynamo.reset() + logits = torch.randn(32, 32, device="cuda", requires_grad=True) + bias = torch.zeros(32, device="cuda") + indices = torch.empty(32, 4, device="cuda", dtype=torch.int32) + histogram = torch.zeros(32, 16, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-2, 2], device="cuda", dtype=torch.float32) + + def router(x): + return fused_topk_with_score_function( + x, + 4, + False, + None, + None, + None, + "sigmoid", + bias, + topk_indices=None if qb else indices, + qb_histogram=histogram if qb else None, + qb_bin_bounds=bounds if qb else None, + qb_histogram_mode="two_kernel" if qb else None, + ) + + with pytest.warns(UserWarning, match="Falling back to eager execution"): + actual = torch.compile(router)(logits) + expected = router(logits) + torch.testing.assert_close(actual, expected) + torch._dynamo.reset() + with pytest.raises(torch._dynamo.exc.Unsupported): + torch.compile(router, fullgraph=True)(logits) + + # Pytorch-based group topk def group_limited_topk( scores: torch.Tensor, diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 898310b0d61..b720565b94e 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -15,11 +15,13 @@ """ import math -from typing import Optional, Union +from typing import Optional, Tuple, Union import torch import transformer_engine_torch as tex +from .jit import no_torch_dynamo + # Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. # Members: # RoutingMapFormat.BYTEMAP — bool[num_tokens, num_experts] @@ -276,6 +278,18 @@ def backward(ctx, grad_probs, _, _histogram_grad): return grad_logits, None, None, None, None, None, None, None, None +def _needs_eager_router(call): + if call.get("topk_indices") is not None: + return "router output buffers" + if any( + call.get(name) is not None + for name in ("qb_histogram", "qb_bin_bounds", "qb_histogram_mode") + ): + return "Quantile Balancing" + return None + + +@no_torch_dynamo(when=_needs_eager_router) def fused_topk_with_score_function( logits: torch.Tensor, topk: int, @@ -293,6 +307,10 @@ def fused_topk_with_score_function( ): """ Fused topk with score function router. + + ``torch.compile`` supports routing-map outputs; ``topk_indices`` and Quantile + Balancing fall back to eager execution and do not support ``fullgraph=True``. + Parameters ---------- logits : torch.Tensor in fp32/bf16/fp16 @@ -375,6 +393,19 @@ def fused_topk_with_score_function( mode, ) return probs, routing_output + if torch.compiler.is_compiling(): + probs, routing_output, _ = _topk_forward( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + ) + return probs, routing_output return FusedTopkScoreFunction.apply( logits, topk, @@ -444,6 +475,7 @@ def fused_compute_score_for_moe_aux_loss( ): """ Fused compute scores for MoE aux loss, subset of the fused_topk_with_score_function. + Supports ``torch.compile(fullgraph=True)`` for forward and backward. Parameters ---------- logits : torch.Tensor in fp32/bf16/fp16 @@ -463,6 +495,9 @@ def fused_compute_score_for_moe_aux_loss( Same shape as ``logits``. """ routing_map_format = _validate_routing_map_format(routing_map_format) + if torch.compiler.is_compiling(): + scores, routing_map, _ = _score_forward(logits, topk, score_function, routing_map_format) + return routing_map, scores return FusedComputeScoresForMoEAuxLoss.apply(logits, topk, score_function, routing_map_format) @@ -537,6 +572,7 @@ def fused_moe_aux_loss( ) -> torch.Tensor: """ Fused MoE aux loss. + Supports ``torch.compile(fullgraph=True)`` with either form of ``total_num_tokens``. Parameters ---------- probs : torch.Tensor in fp32/bf16/fp16 @@ -558,4 +594,306 @@ def fused_moe_aux_loss( aux_loss : torch.Tensor. A scalar tensor in the same dtype as the "probs". """ + if torch.compiler.is_compiling(): + token_count = total_num_tokens if isinstance(total_num_tokens, torch.Tensor) else None + count = 0 if token_count is not None else total_num_tokens + aux_loss, _ = _aux_loss_forward( + probs, tokens_per_expert, count, token_count, num_experts, topk, coeff + ) + return aux_loss return FusedAuxLoss.apply(probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff) + + +def _routing_map_fake(logits, routing_map_format): + if routing_map_format == _ROUTING_MAP_FORMAT_FROM_STRING["bitmap_u8"]: + return logits.new_empty( + (*logits.shape[:-1], (logits.shape[-1] + 7) // 8), dtype=torch.uint8 + ) + return logits.new_empty(logits.shape, dtype=torch.bool) + + +@torch.library.custom_op("te_moe::router_topk_fwd", mutates_args=()) +def _topk_forward( + logits: torch.Tensor, + topk: int, + use_pre_softmax: bool, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, + expert_bias: Optional[torch.Tensor], + routing_map_format: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + with torch.cuda.device(logits.device): + return tex.fused_topk_with_score_function_fwd( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + None, + ) + + +@_topk_forward.register_fake +def _topk_forward_fake( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, +): + del ( + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + ) + return ( + logits.new_empty(logits.shape), + _routing_map_fake(logits, routing_map_format), + logits.new_empty(logits.shape, dtype=torch.float32), + ) + + +@torch.library.custom_op("te_moe::router_topk_bwd", mutates_args=()) +def _topk_backward( + routing_map: torch.Tensor, + intermediate_output: torch.Tensor, + grad_probs: torch.Tensor, + topk: int, + use_pre_softmax: bool, + scaling_factor: Optional[float], + score_function: str, + routing_map_format: int, +) -> torch.Tensor: + grad_probs = grad_probs.contiguous() + grad_logits = torch.empty_like(grad_probs) + with torch.cuda.device(grad_probs.device): + tex.fused_topk_with_score_function_bwd( + routing_map, + intermediate_output, + grad_probs, + grad_logits, + topk, + use_pre_softmax, + scaling_factor, + score_function, + False, + routing_map_format, + ) + return grad_logits + + +@_topk_backward.register_fake +def _topk_backward_fake( + routing_map, + intermediate_output, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, + routing_map_format, +): + del routing_map, intermediate_output, topk, use_pre_softmax, scaling_factor + del score_function, routing_map_format + return grad_probs.new_empty(grad_probs.shape) + + +def _topk_setup_context(ctx, inputs, output): + ( + _, + ctx.topk, + ctx.use_pre_softmax, + _, + _, + ctx.scaling_factor, + ctx.score_function, + _, + ctx.routing_map_format, + ) = inputs + _, routing_map, intermediate_output = output + ctx.mark_non_differentiable(routing_map, intermediate_output) + ctx.save_for_backward(routing_map, intermediate_output) + + +def _topk_autograd(ctx, grad_probs, _grad_map, _grad_intermediate): + routing_map, intermediate_output = ctx.saved_tensors + grad_logits = _topk_backward( + routing_map, + intermediate_output, + grad_probs, + ctx.topk, + ctx.use_pre_softmax, + ctx.scaling_factor, + ctx.score_function, + ctx.routing_map_format, + ) + return grad_logits, None, None, None, None, None, None, None, None + + +_topk_forward.register_autograd(_topk_autograd, setup_context=_topk_setup_context) + + +@torch.library.custom_op("te_moe::router_score_fwd", mutates_args=()) +def _score_forward( + logits: torch.Tensor, + topk: int, + score_function: str, + routing_map_format: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + with torch.cuda.device(logits.device): + return tex.fused_score_for_moe_aux_loss_fwd( + logits, topk, score_function, routing_map_format + ) + + +@_score_forward.register_fake +def _score_forward_fake(logits, topk, score_function, routing_map_format): + del topk, score_function + return ( + logits.new_empty(logits.shape, dtype=torch.float32), + _routing_map_fake(logits, routing_map_format), + logits.new_empty(logits.shape, dtype=torch.float32), + ) + + +@torch.library.custom_op("te_moe::router_score_bwd", mutates_args=()) +def _score_backward( + intermediate_output: torch.Tensor, + grad_scores: torch.Tensor, + topk: int, + score_function: str, + logits_dtype: torch.dtype, +) -> torch.Tensor: + grad_scores = grad_scores.contiguous() + grad_logits = grad_scores.new_empty(grad_scores.shape, dtype=logits_dtype) + with torch.cuda.device(grad_scores.device): + tex.fused_score_for_moe_aux_loss_bwd( + intermediate_output, grad_scores, grad_logits, topk, score_function + ) + return grad_logits + + +@_score_backward.register_fake +def _score_backward_fake(intermediate_output, grad_scores, topk, score_function, logits_dtype): + del intermediate_output, topk, score_function + return grad_scores.new_empty(grad_scores.shape, dtype=logits_dtype) + + +def _score_setup_context(ctx, inputs, output): + logits, ctx.topk, ctx.score_function, _ = inputs + ctx.logits_dtype = logits.dtype + _, routing_map, intermediate_output = output + ctx.mark_non_differentiable(routing_map, intermediate_output) + ctx.save_for_backward(intermediate_output) + + +def _score_autograd(ctx, grad_scores, _grad_map, _grad_intermediate): + (intermediate_output,) = ctx.saved_tensors + grad_logits = _score_backward( + intermediate_output, grad_scores, ctx.topk, ctx.score_function, ctx.logits_dtype + ) + return grad_logits, None, None, None + + +_score_forward.register_autograd(_score_autograd, setup_context=_score_setup_context) + + +@torch.library.custom_op("te_moe::router_aux_loss_fwd", mutates_args=()) +def _aux_loss_forward( + probs: torch.Tensor, + tokens_per_expert: torch.Tensor, + total_num_tokens: int, + token_count: Optional[torch.Tensor], + num_experts: int, + topk: int, + coeff: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + with torch.cuda.device(probs.device): + if token_count is not None: + return tex.fused_moe_aux_loss_fwd_graph_safe( + probs, + tokens_per_expert, + token_count, + num_experts, + probs.shape[0], + probs.shape[1], + topk, + coeff, + ) + return tex.fused_moe_aux_loss_fwd( + probs, + tokens_per_expert, + total_num_tokens, + num_experts, + probs.shape[0], + probs.shape[1], + topk, + coeff, + ) + + +@_aux_loss_forward.register_fake +def _aux_loss_forward_fake( + probs, + tokens_per_expert, + total_num_tokens, + token_count, + num_experts, + topk, + coeff, +): + del tokens_per_expert, total_num_tokens, token_count, num_experts, topk, coeff + return probs.new_empty(()), probs.new_empty((2,), dtype=torch.float32) + + +@torch.library.custom_op("te_moe::router_aux_loss_bwd", mutates_args=()) +def _aux_loss_backward( + const_buf: torch.Tensor, + tokens_per_expert: torch.Tensor, + num_rows: int, + num_cols: int, + grad_aux_loss: torch.Tensor, +) -> torch.Tensor: + with torch.cuda.device(grad_aux_loss.device): + return tex.fused_moe_aux_loss_bwd( + const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss + ) + + +@_aux_loss_backward.register_fake +def _aux_loss_backward_fake(const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss): + del const_buf, tokens_per_expert + return grad_aux_loss.new_empty((num_rows, num_cols)) + + +def _aux_loss_setup_context(ctx, inputs, output): + probs, tokens_per_expert, _, _, _, _, _ = inputs + _, const_buf = output + ctx.num_rows, ctx.num_cols = probs.shape + ctx.mark_non_differentiable(const_buf) + ctx.save_for_backward(const_buf, tokens_per_expert) + + +def _aux_loss_autograd(ctx, grad_aux_loss, _grad_const): + const_buf, tokens_per_expert = ctx.saved_tensors + grad_probs = _aux_loss_backward( + const_buf, tokens_per_expert, ctx.num_rows, ctx.num_cols, grad_aux_loss + ) + return grad_probs, None, None, None, None, None, None + + +_aux_loss_forward.register_autograd(_aux_loss_autograd, setup_context=_aux_loss_setup_context) From 7f30b44e1d8437124c5a234cc15e147630bb4130 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 14 Oct 2026 15:27:38 +0200 Subject: [PATCH 02/14] Support compiled dense router indices without overwriting buffers Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 99 ++++++++++++++++++++++------ transformer_engine/pytorch/router.py | 44 ++++++++++--- 2 files changed, 115 insertions(+), 28 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index f87fee70efa..2d61ad11567 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -133,7 +133,11 @@ def _check_compile_reference(actual, expected, dtype): @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) -@pytest.mark.parametrize("routing_map_format", [RoutingMapFormat.BYTEMAP, "bitmap_u8"]) +@pytest.mark.parametrize( + "routing_map_format,index_dtype", + [(RoutingMapFormat.BYTEMAP, None), ("bitmap_u8", None)] + + [(RoutingMapFormat.BYTEMAP, dtype) for dtype in (torch.int16, torch.int32, torch.int64)], +) @pytest.mark.parametrize( "score_function,use_pre_softmax,grouped", [ @@ -143,7 +147,9 @@ def _check_compile_reference(actual, expected, dtype): ("sqrtsoftplus", False, False), ], ) -def test_topk_compile(dtype, mode, routing_map_format, score_function, use_pre_softmax, grouped): +def test_topk_compile( + dtype, mode, routing_map_format, index_dtype, score_function, use_pre_softmax, grouped +): torch._dynamo.reset() counters.clear() bias = None if score_function == "softmax" else torch.linspace(0.3, 0, 32, device="cuda") @@ -157,12 +163,21 @@ def test_topk_compile(dtype, mode, routing_map_format, score_function, use_pre_s expert_bias=bias, ) - def router(logits): + def router(logits, indices): return fused_topk_with_score_function( - logits, **options, routing_map_format=routing_map_format + logits, + **options, + routing_map_format=routing_map_format, + topk_indices=indices, + overwrite_topk_indices=False, ) compiled = torch.compile(router, fullgraph=True, mode=mode) + indices = ( + torch.full((2, 32, 4), -1, device="cuda", dtype=index_dtype) + if index_dtype is not None + else None + ) for reverse in (False, True, False): torch.compiler.cudagraph_mark_step_begin() logits = _compile_logits(dtype) @@ -170,15 +185,28 @@ def router(logits): logits = logits.flip(-1).detach().requires_grad_(True) eager_logits = logits.detach().clone().requires_grad_(True) ref_logits = logits.detach().float().reshape(-1, 32).requires_grad_(True) - probs, routing_map = compiled(logits) - eager_probs, eager_map = router(eager_logits) + original_indices = indices.clone() if indices is not None else None + original_version = indices._version if indices is not None else None + probs, routing_map = compiled(logits, indices) + eager_probs, eager_map = router(eager_logits, indices) ref_probs, ref_map = topk_score_function_pytorch(ref_logits, **options) ref_map = ref_map.reshape(logits.shape) if routing_map_format == "bitmap_u8": ref_map = _bytemap_to_bitmap_u8(ref_map) assert not routing_map.requires_grad torch.testing.assert_close(routing_map, eager_map) - torch.testing.assert_close(routing_map, ref_map) + if indices is not None: + assert routing_map.shape == indices.shape and routing_map.dtype == index_dtype + assert routing_map.data_ptr() != indices.data_ptr() + assert eager_map.data_ptr() != indices.data_ptr() + assert indices._version == original_version + torch.testing.assert_close(indices, original_indices) + actual_map = topk_indices_to_routing_map(routing_map.reshape(-1, 4), 32) + torch.testing.assert_close(actual_map.reshape(logits.shape), ref_map) + # The caller's buffer must not be a saved backward operand. + indices.fill_(-2 if reverse else -3) + else: + torch.testing.assert_close(routing_map, ref_map) torch.testing.assert_close(probs, eager_probs, atol=0, rtol=0) _check_compile_reference(probs, ref_probs.reshape(logits.shape), dtype) grad = torch.randn((*probs.shape, 2), dtype=dtype, device="cuda")[..., 0] @@ -190,9 +218,12 @@ def router(logits): _check_compile_reference(actual_grad, ref_grad.reshape(logits.shape), dtype) with torch.no_grad(): torch.compiler.cudagraph_mark_step_begin() - actual = compiled(logits) - expected = router(logits) + actual = compiled(logits, indices) + expected = router(logits, indices) torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if indices is not None: + assert actual[1].data_ptr() != indices.data_ptr() + torch.testing.assert_close(indices, torch.full_like(indices, -3)) if mode == "reduce-overhead": assert not counters["inductor"]["cudagraph_skips"] @@ -285,13 +316,29 @@ def loss(probs, total_num_tokens): @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -def test_router_compile_dynamic(dtype): +@pytest.mark.parametrize("index_dtype", [None, torch.int16, torch.int32, torch.int64]) +def test_router_compile_dynamic(dtype, index_dtype): torch._dynamo.reset() counts = torch.arange(1, 131, device="cuda", dtype=torch.int32) def router_loss(logits): + indices = ( + torch.empty(logits.shape[0], 4, dtype=index_dtype, device=logits.device) + if index_dtype is not None + else None + ) probs, routing_map = fused_topk_with_score_function( - logits, 4, False, None, None, None, "sigmoid", None, "bitmap_u8" + logits, + 4, + False, + None, + None, + None, + "sigmoid", + None, + "bytemap" if indices is not None else "bitmap_u8", + topk_indices=indices, + overwrite_topk_indices=False, ) _, scores = fused_compute_score_for_moe_aux_loss(logits, 4, "sigmoid", 1) aux_loss = fused_moe_aux_loss(scores, counts, logits.shape[0], 130, 4, 0.01) @@ -303,7 +350,7 @@ def router_loss(logits): eager_logits = logits.detach().clone().requires_grad_(True) actual, routing_map = compiled(logits) expected, eager_map = router_loss(eager_logits) - assert routing_map.shape == (num_tokens, 17) + assert routing_map.shape == (num_tokens, 4 if index_dtype is not None else 17) torch.testing.assert_close(routing_map, eager_map) torch.testing.assert_close(actual, expected) (actual_grad,) = torch.autograd.grad(actual, logits) @@ -312,14 +359,22 @@ def router_loss(logits): @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("routing_map_format", [0, 1]) -def test_router_compile_opcheck(dtype, routing_map_format): +@pytest.mark.parametrize( + "routing_map_format,index_dtype", + [(0, None), (1, None)] + [(0, dtype) for dtype in (torch.int16, torch.int32, torch.int64)], +) +def test_router_compile_opcheck(dtype, routing_map_format, index_dtype): from transformer_engine.pytorch import router logits = torch.randn(16, 130, device="cuda", dtype=dtype, requires_grad=True) + indices = ( + torch.full((16, 4), -1, device="cuda", dtype=index_dtype) + if index_dtype is not None + else None + ) torch.library.opcheck( router._topk_forward, - (logits, 4, False, None, None, None, "sigmoid", None, routing_map_format), + (logits, 4, False, None, None, None, "sigmoid", None, routing_map_format, indices), ) torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) counts = torch.ones(130, device="cuda", dtype=torch.int32) @@ -330,8 +385,8 @@ def test_router_compile_opcheck(dtype, routing_map_format): ) -@pytest.mark.parametrize("qb", [False, True]) -def test_router_compile_eager_fallback(qb): +@pytest.mark.parametrize("qb,overwrite", [(False, True), (True, True), (True, False)]) +def test_router_compile_eager_fallback(qb, overwrite): torch._dynamo.reset() logits = torch.randn(32, 32, device="cuda", requires_grad=True) bias = torch.zeros(32, device="cuda") @@ -353,6 +408,7 @@ def router(x): qb_histogram=histogram if qb else None, qb_bin_bounds=bounds if qb else None, qb_histogram_mode="two_kernel" if qb else None, + **({} if overwrite else {"overwrite_topk_indices": False}), ) with pytest.warns(UserWarning, match="Falling back to eager execution"): @@ -740,7 +796,8 @@ def test_topk_sqrtsoftplus( "routing_output_mode", ["bytemap", "bitmap_u8", "dense_int16", "dense_int32", "dense_int64"], ) -def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode): +@pytest.mark.parametrize("overwrite", [False, True]) +def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, overwrite): num_tokens = 257 num_experts = 896 num_bins = 1000 @@ -767,7 +824,7 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode): "dense_int64": torch.int64, }.get(routing_output_mode) topk_indices = ( - torch.empty(num_tokens, topk, device="cuda", dtype=dense_dtype) + torch.full((num_tokens, topk), -1, device="cuda", dtype=dense_dtype) if dense_dtype is not None else None ) @@ -790,9 +847,13 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode): qb_histogram=fused_histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, + overwrite_topk_indices=overwrite, ) torch.testing.assert_close(fused_probs, reference["probs"]) if dense_dtype is not None: + assert (fused_routing_output.data_ptr() == topk_indices.data_ptr()) == overwrite + if not overwrite: + torch.testing.assert_close(topk_indices, torch.full_like(topk_indices, -1)) fused_routing_map = topk_indices_to_routing_map(fused_routing_output, num_experts) torch.testing.assert_close(fused_routing_map, reference["routing_map"]) elif routing_output_mode == "bitmap_u8": diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index b720565b94e..c2a59c61af0 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -279,7 +279,7 @@ def backward(ctx, grad_probs, _, _histogram_grad): def _needs_eager_router(call): - if call.get("topk_indices") is not None: + if call.get("topk_indices") is not None and call.get("overwrite_topk_indices", True): return "router output buffers" if any( call.get(name) is not None @@ -304,12 +304,14 @@ def fused_topk_with_score_function( qb_histogram: Optional[torch.Tensor] = None, qb_bin_bounds: Optional[torch.Tensor] = None, qb_histogram_mode: Optional[str] = None, + overwrite_topk_indices: bool = True, ): """ Fused topk with score function router. - ``torch.compile`` supports routing-map outputs; ``topk_indices`` and Quantile - Balancing fall back to eager execution and do not support ``fullgraph=True``. + ``torch.compile`` supports routing maps and dense indices with + ``overwrite_topk_indices=False``. Overwriting index buffers and Quantile Balancing + fall back to eager execution and do not support ``fullgraph=True``. Parameters ---------- @@ -333,8 +335,9 @@ def fused_topk_with_score_function( bit ``(e % 8)`` of byte ``(e / 8)`` set when token ``t`` routes to expert ``e`` (LSB-first / little-endian packing along the expert axis). topk_indices : torch.Tensor, optional - Optional output buffer with shape [num_tokens, topk]. When provided, its dtype + Optional output buffer with shape [*logits.shape[:-1], topk]. When provided, its dtype controls the dense index output dtype and the routing map is not materialized. + With ``overwrite_topk_indices=False``, only its metadata is used. qb_histogram : torch.Tensor, optional Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. qb_bin_bounds : torch.Tensor, optional @@ -348,6 +351,9 @@ def fused_topk_with_score_function( Invalid replay-time bounds may silently produce an incorrect histogram. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. + overwrite_topk_indices : bool, default = True + Write indices into ``topk_indices`` and return that buffer. If False, return a new + tensor without modifying the supplied buffer. Ignored when ``topk_indices`` is None. Returns ------- @@ -359,10 +365,17 @@ def fused_topk_with_score_function( - BYTEMAP: bool[*logits.shape[:-1], num_experts] - BITMAP_U8: uint8[*logits.shape[:-1], ceil(num_experts/8)] LSB-first bit-packed. + - Dense indices: [*logits.shape[:-1], topk], with the dtype of ``topk_indices``. """ if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type.") routing_map_format = _validate_routing_map_format(routing_map_format) + if ( + not torch.compiler.is_compiling() + and topk_indices is not None + and not overwrite_topk_indices + ): + topk_indices = torch.empty_like(topk_indices) qb_arguments = (qb_histogram, qb_bin_bounds, qb_histogram_mode) if any(value is not None for value in qb_arguments): if any(value is None for value in qb_arguments): @@ -404,6 +417,7 @@ def fused_topk_with_score_function( score_function, expert_bias, routing_map_format, + topk_indices, ) return probs, routing_output return FusedTopkScoreFunction.apply( @@ -623,8 +637,10 @@ def _topk_forward( score_function: str, expert_bias: Optional[torch.Tensor], routing_map_format: int, + topk_indices: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: with torch.cuda.device(logits.device): + indices = torch.empty_like(topk_indices) if topk_indices is not None else None return tex.fused_topk_with_score_function_fwd( logits, topk, @@ -635,7 +651,7 @@ def _topk_forward( score_function, expert_bias, routing_map_format, - None, + indices, ) @@ -650,6 +666,7 @@ def _topk_forward_fake( score_function, expert_bias, routing_map_format, + topk_indices=None, ): del ( topk, @@ -662,7 +679,11 @@ def _topk_forward_fake( ) return ( logits.new_empty(logits.shape), - _routing_map_fake(logits, routing_map_format), + ( + torch.empty_like(topk_indices) + if topk_indices is not None + else _routing_map_fake(logits, routing_map_format) + ), logits.new_empty(logits.shape, dtype=torch.float32), ) @@ -677,6 +698,7 @@ def _topk_backward( scaling_factor: Optional[float], score_function: str, routing_map_format: int, + use_dense_indices: bool = False, ) -> torch.Tensor: grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -690,7 +712,7 @@ def _topk_backward( use_pre_softmax, scaling_factor, score_function, - False, + use_dense_indices, routing_map_format, ) return grad_logits @@ -706,9 +728,10 @@ def _topk_backward_fake( scaling_factor, score_function, routing_map_format, + use_dense_indices=False, ): del routing_map, intermediate_output, topk, use_pre_softmax, scaling_factor - del score_function, routing_map_format + del score_function, routing_map_format, use_dense_indices return grad_probs.new_empty(grad_probs.shape) @@ -723,7 +746,9 @@ def _topk_setup_context(ctx, inputs, output): ctx.score_function, _, ctx.routing_map_format, + topk_indices, ) = inputs + ctx.use_dense_indices = topk_indices is not None _, routing_map, intermediate_output = output ctx.mark_non_differentiable(routing_map, intermediate_output) ctx.save_for_backward(routing_map, intermediate_output) @@ -740,8 +765,9 @@ def _topk_autograd(ctx, grad_probs, _grad_map, _grad_intermediate): ctx.scaling_factor, ctx.score_function, ctx.routing_map_format, + ctx.use_dense_indices, ) - return grad_logits, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None _topk_forward.register_autograd(_topk_autograd, setup_context=_topk_setup_context) From bd8328c70cb5cc403cb3c8bb3d783019688e7c33 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 14 Oct 2026 15:52:26 +0200 Subject: [PATCH 03/14] Support compiled quantile balancing with functional histogram outputs Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 305 ++++++++++++++++++++++++++- transformer_engine/pytorch/router.py | 144 ++++++++++++- 2 files changed, 437 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 2d61ad11567..6f1addea2f9 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -385,8 +385,18 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype): ) -@pytest.mark.parametrize("qb,overwrite", [(False, True), (True, True), (True, False)]) -def test_router_compile_eager_fallback(qb, overwrite): +@pytest.mark.parametrize( + "qb,overwrite,delta,assume_valid", + [ + (False, True, False, False), + (True, True, False, False), + (True, False, False, False), + (True, False, False, True), + (True, False, True, False), + (True, True, True, True), + ], +) +def test_router_compile_eager_fallback(qb, overwrite, delta, assume_valid): torch._dynamo.reset() logits = torch.randn(32, 32, device="cuda", requires_grad=True) bias = torch.zeros(32, device="cuda") @@ -404,11 +414,13 @@ def router(x): None, "sigmoid", bias, - topk_indices=None if qb else indices, + topk_indices=indices if not qb or (delta and assume_valid) else None, qb_histogram=histogram if qb else None, qb_bin_bounds=bounds if qb else None, qb_histogram_mode="two_kernel" if qb else None, **({} if overwrite else {"overwrite_topk_indices": False}), + return_qb_histogram_delta=delta, + assume_valid_qb_bin_bounds=assume_valid, ) with pytest.warns(UserWarning, match="Falling back to eager execution"): @@ -420,6 +432,289 @@ def router(x): torch.compile(router, fullgraph=True)(logits) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +@pytest.mark.parametrize( + "routing_map_format,index_dtype", + [("bytemap", None), ("bitmap_u8", None)] + + [("bytemap", dtype) for dtype in (torch.int16, torch.int32, torch.int64)], +) +@pytest.mark.parametrize("topk", [1, 4, 16]) +def test_qb_compile(dtype, mode, histogram_mode, routing_map_format, index_dtype, topk): + torch._dynamo.reset() + counters.clear() + bias = torch.linspace(0.3, 0, 32, device="cuda") + histogram = torch.full((32, 31), 7, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.1, 1.2], device="cuda") + indices = ( + torch.full((2, 32, topk), -1, device="cuda", dtype=index_dtype) + if index_dtype is not None + else None + ) + + def router(logits, indices, histogram, bounds): + return fused_topk_with_score_function( + logits, + topk, + False, + None, + None, + 1.2, + "sigmoid", + bias, + routing_map_format=routing_map_format, + topk_indices=indices, + qb_histogram=histogram, + qb_bin_bounds=bounds, + qb_histogram_mode=histogram_mode, + overwrite_topk_indices=False, + return_qb_histogram_delta=True, + assume_valid_qb_bin_bounds=True, + ) + + compiled = torch.compile(router, fullgraph=True, mode=mode) + accumulated = torch.zeros_like(histogram) + reference_accumulated = torch.zeros_like(histogram) + for step, new_bounds in enumerate(([-1.1, 1.2], [-0.3, 0.7], [-0.8, 1.4])): + torch.compiler.cudagraph_mark_step_begin() + bounds.copy_(bounds.new_tensor(new_bounds)) + logits = _compile_logits(dtype) + if step == 1: + logits = logits.flip(-1).detach().requires_grad_(True) + eager_logits = logits.detach().clone().requires_grad_(True) + ref_logits = logits.detach().float().requires_grad_(True) + histogram_version = histogram._version + indices_version = indices._version if indices is not None else None + probs, routing, delta = compiled(logits, indices, histogram, bounds) + eager_probs, eager_routing, eager_delta = router(eager_logits, indices, histogram, bounds) + reference = qb_topk_score_function_pytorch(ref_logits, topk, bias, bounds, 31) + ref_probs = reference["probs"] * 1.2 + ref_map = reference["routing_map"] + assert not routing.requires_grad and not delta.requires_grad + assert delta.data_ptr() != histogram.data_ptr() + assert eager_delta.data_ptr() != histogram.data_ptr() + assert histogram._version == histogram_version + torch.testing.assert_close(histogram, torch.full_like(histogram, 7 + step)) + torch.testing.assert_close(delta, eager_delta) + torch.testing.assert_close(delta, reference["histogram"]) + torch.testing.assert_close(routing, eager_routing) + if indices is not None: + assert routing.shape == indices.shape and routing.dtype == index_dtype + assert routing.data_ptr() != indices.data_ptr() + assert eager_routing.data_ptr() != indices.data_ptr() + assert indices._version == indices_version + torch.testing.assert_close(indices, torch.full_like(indices, -1 - step)) + actual_map = topk_indices_to_routing_map(routing.reshape(-1, topk), 32) + torch.testing.assert_close(actual_map.reshape(logits.shape), ref_map) + indices.sub_(1) + elif routing_map_format == "bitmap_u8": + torch.testing.assert_close(routing, _bytemap_to_bitmap_u8(ref_map)) + else: + torch.testing.assert_close(routing, ref_map) + accumulated.add_(delta) + reference_accumulated.add_(reference["histogram"]) + torch.testing.assert_close(accumulated, reference_accumulated) + torch.testing.assert_close(probs, eager_probs, atol=0, rtol=0) + _check_compile_reference(probs, ref_probs, dtype) + histogram.add_(1) + bounds.add_(0.1) + grad = torch.randn((*probs.shape, 2), dtype=dtype, device="cuda")[..., 0] + (actual_grad,) = torch.autograd.grad(probs, logits, grad) + (eager_grad,) = torch.autograd.grad(eager_probs, eager_logits, grad) + (ref_grad,) = torch.autograd.grad(ref_probs, ref_logits, grad.float()) + torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) + _check_compile_reference(actual_grad, ref_grad, dtype) + with torch.no_grad(): + torch.compiler.cudagraph_mark_step_begin() + torch.testing.assert_close( + compiled(logits, indices, histogram, bounds), + router(logits, indices, histogram, bounds), + atol=0, + rtol=0, + ) + if mode == "reduce-overhead": + assert not counters["inductor"]["cudagraph_skips"] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_compile_dynamic_accumulation(dtype, histogram_mode): + torch._dynamo.reset() + bias = torch.linspace(0.3, 0, 130, device="cuda") + bounds = torch.tensor([-1.0, 1.0], device="cuda") + histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) + reference_histogram = histogram.clone() + + def router(logits, histogram, bounds): + probs, routing, delta = fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + bias, + "bitmap_u8", + qb_histogram=histogram, + qb_bin_bounds=bounds, + qb_histogram_mode=histogram_mode, + return_qb_histogram_delta=True, + assume_valid_qb_bin_bounds=True, + ) + return probs, routing, histogram + delta + + compiled = torch.compile(router, fullgraph=True, dynamic=True) + for num_tokens in (17, 33, 64): + logits = torch.randn(num_tokens, 130, device="cuda", dtype=dtype, requires_grad=True) + input_histogram = histogram.clone() + probs, routing, histogram = compiled(logits, input_histogram, bounds) + reference = qb_topk_score_function_pytorch(logits, 4, bias, bounds, 64, reference_histogram) + torch.testing.assert_close(histogram, reference["histogram"]) + torch.testing.assert_close(routing, _bytemap_to_bitmap_u8(reference["routing_map"])) + _check_compile_reference(probs, reference["probs"], dtype) + assert histogram.data_ptr() != input_histogram.data_ptr() + torch.autograd.grad(probs.sum(), logits) + bounds.add_(0.03) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "histogram_mode", [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC] +) +@pytest.mark.parametrize( + "routing_map_format,index_dtype", + [(0, None), (1, None)] + [(0, dtype) for dtype in (torch.int16, torch.int32, torch.int64)], +) +def test_qb_compile_opcheck(dtype, histogram_mode, routing_map_format, index_dtype): + from transformer_engine.pytorch import router + + logits = torch.randn(16, 130, device="cuda", dtype=dtype, requires_grad=True) + bias = torch.zeros(130, device="cuda") + histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.0, 1.0], device="cuda") + indices = ( + torch.full((16, 4), -1, device="cuda", dtype=index_dtype) + if index_dtype is not None + else None + ) + torch.library.opcheck( + router._qb_forward, + ( + logits, + 4, + None, + bias, + routing_map_format, + indices, + histogram, + bounds, + int(histogram_mode), + ), + ) + + +@pytest.mark.parametrize("delta", [False, True]) +def test_qb_trusted_bounds_skip_value_validation(monkeypatch, delta): + from transformer_engine.pytorch import router + + logits = torch.randn(8, 16, device="cuda", requires_grad=True) + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.0, 1.0], device="cuda") + + def unexpected_validation(_): + raise AssertionError("trusted QB bounds must not read values on the CPU") + + with monkeypatch.context() as patch: + patch.setattr(router, "_validate_qb_bin_bounds", unexpected_validation) + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + torch.zeros(16, device="cuda"), + qb_histogram=histogram, + qb_bin_bounds=bounds, + qb_histogram_mode="two_kernel", + return_qb_histogram_delta=delta, + assume_valid_qb_bin_bounds=True, + ) + bounds.zero_() + with pytest.raises(ValueError, match="finite with lower < upper"): + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + torch.zeros(16, device="cuda"), + qb_histogram=histogram, + qb_bin_bounds=bounds, + qb_histogram_mode="two_kernel", + return_qb_histogram_delta=delta, + ) + + +@pytest.mark.parametrize("compiled", [False, True]) +@pytest.mark.parametrize("invalid", ["shape", "dtype", "device"]) +def test_qb_trusted_bounds_check_metadata(compiled, invalid): + torch._dynamo.reset() + logits = torch.randn(8, 16, device="cuda") + bias = torch.zeros(16, device="cuda") + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.0, 1.0], device="cuda") + if invalid == "shape": + bounds = bounds.reshape(1, 2) + elif invalid == "dtype": + bounds = bounds.half() + else: + bounds = bounds.cpu() + + def router(logits, bounds): + return fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + bias, + qb_histogram=histogram, + qb_bin_bounds=bounds, + qb_histogram_mode="two_kernel", + return_qb_histogram_delta=True, + assume_valid_qb_bin_bounds=True, + ) + + if compiled: + router = torch.compile(router, fullgraph=True) + with pytest.raises(RuntimeError, match="QB bin_bounds"): + router(logits, bounds) + + +def test_qb_histogram_delta_requires_qb(): + logits = torch.randn(8, 16, device="cuda") + with pytest.raises(ValueError, match="return_qb_histogram_delta requires"): + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + None, + return_qb_histogram_delta=True, + ) + + # Pytorch-based group topk def group_limited_topk( scores: torch.Tensor, @@ -976,7 +1271,8 @@ def test_qb_topk_uses_logits_device(histogram_mode): @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @pytest.mark.parametrize("invalid_bounds", ["equal", "reversed", "nonfinite"]) -def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds): +@pytest.mark.parametrize("delta", [False, True]) +def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds, delta): logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) @@ -999,6 +1295,7 @@ def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds): qb_histogram=histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, + return_qb_histogram_delta=delta, ) diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index c2a59c61af0..7d0e216f121 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -223,9 +223,10 @@ def forward( histogram: torch.Tensor, bin_bounds: torch.Tensor, histogram_mode: int, + assume_valid_bin_bounds: bool, ): # pylint: disable=missing-function-docstring - bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) + bin_bounds_validated = assume_valid_bin_bounds or _validate_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -275,7 +276,7 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.use_dense_indices, ctx.routing_map_format, ) - return grad_logits, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None def _needs_eager_router(call): @@ -285,7 +286,10 @@ def _needs_eager_router(call): call.get(name) is not None for name in ("qb_histogram", "qb_bin_bounds", "qb_histogram_mode") ): - return "Quantile Balancing" + if not call.get("return_qb_histogram_delta", False): + return "in-place QB histogram accumulation" + if not call.get("assume_valid_qb_bin_bounds", False): + return "QB bin-bound value validation" return None @@ -305,13 +309,17 @@ def fused_topk_with_score_function( qb_bin_bounds: Optional[torch.Tensor] = None, qb_histogram_mode: Optional[str] = None, overwrite_topk_indices: bool = True, + return_qb_histogram_delta: bool = False, + assume_valid_qb_bin_bounds: bool = False, ): """ Fused topk with score function router. ``torch.compile`` supports routing maps and dense indices with - ``overwrite_topk_indices=False``. Overwriting index buffers and Quantile Balancing - fall back to eager execution and do not support ``fullgraph=True``. + ``overwrite_topk_indices=False``. Quantile Balancing also requires + ``return_qb_histogram_delta=True`` and ``assume_valid_qb_bin_bounds=True``. + Buffer mutation and QB value validation fall back to eager execution and do not + support ``fullgraph=True``. Parameters ---------- @@ -340,10 +348,13 @@ def fused_topk_with_score_function( With ``overwrite_topk_indices=False``, only its metadata is used. qb_histogram : torch.Tensor, optional Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. + With ``return_qb_histogram_delta=True``, only its metadata is used; the supplied + histogram is neither read nor modified. qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be - finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. - Before CUDA graph capture, validate the current version eagerly or use + finite with ``lower < upper``. With default value validation, eager calls revalidate + PyTorch-tracked in-place updates. Before eager CUDA graph capture, validate the + current version eagerly or use :func:`mark_qb_bin_bounds_validated` after a trusted device-side update. CUDA graph replay does not execute Python validation or advance PyTorch tensor-version counters. Any bounds update visible to replay, whether issued outside the graph or captured inside it, is @@ -354,6 +365,15 @@ def fused_topk_with_score_function( overwrite_topk_indices : bool, default = True Write indices into ``topk_indices`` and return that buffer. If False, return a new tensor without modifying the supplied buffer. Ignored when ``topk_indices`` is None. + return_qb_histogram_delta : bool, default = False + Return the current call's histogram as a third output instead of accumulating into + ``qb_histogram``. Requires all three QB arguments. The caller owns accumulation, + for example ``histogram = histogram + histogram_delta`` across microbatches. + assume_valid_qb_bin_bounds : bool, default = False + Skip QB bounds value validation and its CPU synchronization. The caller must + guarantee finite bounds with ``lower < upper`` on every call and CUDA graph replay, + including after device-side updates. Tensor metadata is still checked. Ignored + without QB. Setting this does not mark the tensor as validated for later checked calls. Returns ------- @@ -366,6 +386,10 @@ def fused_topk_with_score_function( - BITMAP_U8: uint8[*logits.shape[:-1], ceil(num_experts/8)] LSB-first bit-packed. - Dense indices: [*logits.shape[:-1], topk], with the dtype of ``topk_indices``. + histogram_delta : torch.Tensor, optional + Returned only when ``return_qb_histogram_delta=True``. A new non-differentiable + int32 ``[num_experts, num_bins]`` tensor containing this call's counts, excluding + all previous counts in ``qb_histogram``. """ if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type.") @@ -394,7 +418,22 @@ def fused_topk_with_score_function( "qb_histogram_mode must be 'two_kernel' or 'fused_atomic', " f"got {qb_histogram_mode!r}" ) - probs, routing_output, _ = FusedTopkScoreFunctionQB.apply( + if torch.compiler.is_compiling(): + probs, routing_output, _, histogram_delta = _qb_forward( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + qb_histogram, + qb_bin_bounds, + mode, + ) + return probs, routing_output, histogram_delta + if return_qb_histogram_delta: + qb_histogram = torch.zeros_like(qb_histogram) + probs, routing_output, histogram_output = FusedTopkScoreFunctionQB.apply( logits, topk, scaling_factor, @@ -404,8 +443,13 @@ def fused_topk_with_score_function( qb_histogram, qb_bin_bounds, mode, + assume_valid_qb_bin_bounds, ) + if return_qb_histogram_delta: + return probs, routing_output, histogram_output return probs, routing_output + if return_qb_histogram_delta: + raise ValueError("return_qb_histogram_delta requires all three QB arguments") if torch.compiler.is_compiling(): probs, routing_output, _ = _topk_forward( logits, @@ -773,6 +817,90 @@ def _topk_autograd(ctx, grad_probs, _grad_map, _grad_intermediate): _topk_forward.register_autograd(_topk_autograd, setup_context=_topk_setup_context) +@torch.library.custom_op("te_moe::router_qb_fwd", mutates_args=()) +def _qb_forward( + logits: torch.Tensor, + topk: int, + scaling_factor: Optional[float], + expert_bias: torch.Tensor, + routing_map_format: int, + topk_indices: Optional[torch.Tensor], + histogram: torch.Tensor, + bin_bounds: torch.Tensor, + histogram_mode: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + with torch.cuda.device(logits.device): + indices = torch.empty_like(topk_indices) if topk_indices is not None else None + histogram_delta = torch.zeros_like(histogram) + probs, routing_output, intermediate_output, _, histogram_delta = ( + tex.fused_topk_with_score_function_qb_fwd( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + indices, + histogram_delta, + bin_bounds, + histogram_mode, + True, + ) + ) + return probs, routing_output, intermediate_output, histogram_delta + + +@_qb_forward.register_fake +def _qb_forward_fake( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, +): + del topk, scaling_factor, expert_bias, bin_bounds, histogram_mode + return ( + logits.new_empty(logits.shape), + ( + torch.empty_like(topk_indices) + if topk_indices is not None + else _routing_map_fake(logits, routing_map_format) + ), + logits.new_empty(logits.shape, dtype=torch.float32), + torch.empty_like(histogram), + ) + + +def _qb_setup_context(ctx, inputs, output): + _, ctx.topk, ctx.scaling_factor, _, ctx.routing_map_format, indices, _, _, _ = inputs + ctx.use_dense_indices = indices is not None + _, routing_output, intermediate_output, histogram_delta = output + ctx.mark_non_differentiable(routing_output, intermediate_output, histogram_delta) + ctx.save_for_backward(routing_output, intermediate_output) + + +def _qb_autograd(ctx, grad_probs, _grad_routing, _grad_intermediate, _grad_histogram): + routing_output, intermediate_output = ctx.saved_tensors + grad_logits = _topk_backward( + routing_output, + intermediate_output, + grad_probs, + ctx.topk, + False, + ctx.scaling_factor, + "sigmoid", + ctx.routing_map_format, + ctx.use_dense_indices, + ) + return grad_logits, None, None, None, None, None, None, None, None + + +_qb_forward.register_autograd(_qb_autograd, setup_context=_qb_setup_context) + + @torch.library.custom_op("te_moe::router_score_fwd", mutates_args=()) def _score_forward( logits: torch.Tensor, From f5c02da6441e3460c9eef70a0c81891f2b98604b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 14:08:20 +0200 Subject: [PATCH 04/14] Reuse router numerical tests for selected torch.compile configurations Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 993 +++++++++++++---------------- 1 file changed, 461 insertions(+), 532 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 6f1addea2f9..7ea10af2f8d 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. from copy import deepcopy +from itertools import product from typing import Optional import pytest @@ -104,9 +105,10 @@ def _get_tolerances(dtype: torch.dtype, num_experts: int): """ # Default tolerances for torch.testing.assert_close base_atol, base_rtol = 1e-5, 1.3e-6 - # TODO: account for fp16, bf16 as dtype - if dtype != torch.float32: - raise NotImplementedError("tolerances implemented for fp32 only") + if dtype == torch.float16: + return 2e-4, 2e-3 + if dtype == torch.bfloat16: + return 2e-3, 2e-2 eps = 2e-7 # The worst-case rounding error from summing N values is O(N * eps). # Use 2 * num_experts * eps as the tolerance floor so tests pass for @@ -116,6 +118,13 @@ def _get_tolerances(dtype: torch.dtype, num_experts: int): return atol, rtol +def _router_function(fn, execution): + if execution == "eager": + return fn + torch._dynamo.reset() + return torch.compile(fn, fullgraph=True) + + def _compile_logits(dtype): logits = torch.linspace(-2, 2, 32, device="cuda") logits = logits + torch.linspace(-0.1, 0.1, 64, device="cuda").unsqueeze(1) @@ -131,192 +140,78 @@ def _check_compile_reference(actual, expected, dtype): torch.testing.assert_close(actual.float(), expected.float(), atol=atol, rtol=rtol) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) -@pytest.mark.parametrize( - "routing_map_format,index_dtype", - [(RoutingMapFormat.BYTEMAP, None), ("bitmap_u8", None)] - + [(RoutingMapFormat.BYTEMAP, dtype) for dtype in (torch.int16, torch.int32, torch.int64)], -) @pytest.mark.parametrize( - "score_function,use_pre_softmax,grouped", - [ - ("softmax", True, False), - ("softmax", False, True), - ("sigmoid", False, True), - ("sqrtsoftplus", False, False), - ], + "histogram_mode,index_dtype", + [(None, torch.int32), ("two_kernel", None), ("fused_atomic", torch.int16)], ) -def test_topk_compile( - dtype, mode, routing_map_format, index_dtype, score_function, use_pre_softmax, grouped -): +def test_router_compile_replay(histogram_mode, index_dtype): torch._dynamo.reset() counters.clear() - bias = None if score_function == "softmax" else torch.linspace(0.3, 0, 32, device="cuda") - options = dict( - topk=4, - use_pre_softmax=use_pre_softmax, - num_groups=4 if grouped else None, - group_topk=2 if grouped else None, - scaling_factor=1.2, - score_function=score_function, - expert_bias=bias, + histogram = torch.full((32, 31), 7, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.1, 1.2], device="cuda") + bias = torch.linspace(0.3, 0, 32, device="cuda") + counts = torch.arange(1, 33, device="cuda", dtype=torch.int32) + total = torch.tensor(64, device="cuda", dtype=torch.int64) + indices = ( + torch.full((2, 32, 4), -1, device="cuda", dtype=index_dtype) + if index_dtype is not None + else None ) - def router(logits, indices): - return fused_topk_with_score_function( + def router(logits, bounds, total): + outputs = fused_topk_with_score_function( logits, - **options, - routing_map_format=routing_map_format, + 4, + False, + None, + None, + 1.2, + "sigmoid", + bias, + routing_map_format="bitmap_u8" if indices is None else "bytemap", topk_indices=indices, overwrite_topk_indices=False, + qb_histogram=histogram if histogram_mode else None, + qb_bin_bounds=bounds if histogram_mode else None, + qb_histogram_mode=histogram_mode, + return_qb_histogram_delta=histogram_mode is not None, + assume_valid_qb_bin_bounds=histogram_mode is not None, ) + _, scores = fused_compute_score_for_moe_aux_loss(logits, 4, "sigmoid") + loss = fused_moe_aux_loss(scores.reshape(-1, 32), counts, total, 32, 4, 0.01) + return (outputs[0].float().square().sum() + loss, *outputs[1:]) - compiled = torch.compile(router, fullgraph=True, mode=mode) - indices = ( - torch.full((2, 32, 4), -1, device="cuda", dtype=index_dtype) - if index_dtype is not None - else None - ) - for reverse in (False, True, False): + compiled = torch.compile(router, fullgraph=True, mode="reduce-overhead") + for step, new_bounds in enumerate(([-1.1, 1.2], [-0.3, 0.7], [-0.8, 1.4])): torch.compiler.cudagraph_mark_step_begin() - logits = _compile_logits(dtype) - if reverse: + bounds.copy_(bounds.new_tensor(new_bounds)) + total.fill_(64 + 17 * step) + logits = _compile_logits(torch.bfloat16) + if step == 1: logits = logits.flip(-1).detach().requires_grad_(True) eager_logits = logits.detach().clone().requires_grad_(True) - ref_logits = logits.detach().float().reshape(-1, 32).requires_grad_(True) - original_indices = indices.clone() if indices is not None else None - original_version = indices._version if indices is not None else None - probs, routing_map = compiled(logits, indices) - eager_probs, eager_map = router(eager_logits, indices) - ref_probs, ref_map = topk_score_function_pytorch(ref_logits, **options) - ref_map = ref_map.reshape(logits.shape) - if routing_map_format == "bitmap_u8": - ref_map = _bytemap_to_bitmap_u8(ref_map) - assert not routing_map.requires_grad - torch.testing.assert_close(routing_map, eager_map) - if indices is not None: - assert routing_map.shape == indices.shape and routing_map.dtype == index_dtype - assert routing_map.data_ptr() != indices.data_ptr() - assert eager_map.data_ptr() != indices.data_ptr() - assert indices._version == original_version - torch.testing.assert_close(indices, original_indices) - actual_map = topk_indices_to_routing_map(routing_map.reshape(-1, 4), 32) - torch.testing.assert_close(actual_map.reshape(logits.shape), ref_map) - # The caller's buffer must not be a saved backward operand. - indices.fill_(-2 if reverse else -3) - else: - torch.testing.assert_close(routing_map, ref_map) - torch.testing.assert_close(probs, eager_probs, atol=0, rtol=0) - _check_compile_reference(probs, ref_probs.reshape(logits.shape), dtype) - grad = torch.randn((*probs.shape, 2), dtype=dtype, device="cuda")[..., 0] - assert not grad.is_contiguous() - (actual_grad,) = torch.autograd.grad(probs, logits, grad) - (eager_grad,) = torch.autograd.grad(eager_probs, eager_logits, grad) - (ref_grad,) = torch.autograd.grad(ref_probs, ref_logits, grad.float().reshape(-1, 32)) - torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) - _check_compile_reference(actual_grad, ref_grad.reshape(logits.shape), dtype) - with torch.no_grad(): - torch.compiler.cudagraph_mark_step_begin() - actual = compiled(logits, indices) - expected = router(logits, indices) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) + actual = compiled(logits, bounds, total) + expected = router(eager_logits, bounds, total) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(histogram, torch.full_like(histogram, 7 + step)) if indices is not None: - assert actual[1].data_ptr() != indices.data_ptr() - torch.testing.assert_close(indices, torch.full_like(indices, -3)) - if mode == "reduce-overhead": - assert not counters["inductor"]["cudagraph_skips"] - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) -@pytest.mark.parametrize("routing_map_format", [RoutingMapFormat.BYTEMAP, "bitmap_u8"]) -@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) -def test_scores_for_aux_loss_compile(dtype, mode, routing_map_format, score_function): - torch._dynamo.reset() - counters.clear() - - def scores(logits): - return fused_compute_score_for_moe_aux_loss(logits, 4, score_function, routing_map_format) - - compiled = torch.compile(scores, fullgraph=True, mode=mode) - for reverse in (False, True, False): - torch.compiler.cudagraph_mark_step_begin() - logits = _compile_logits(dtype) - if reverse: - logits = logits.flip(-1).detach().requires_grad_(True) - eager_logits = logits.detach().clone().requires_grad_(True) - ref_logits = logits.detach().float().reshape(-1, 32).requires_grad_(True) - routing_map, output = compiled(logits) - eager_map, eager_output = scores(eager_logits) - ref_map, ref_output = compute_scores_for_aux_loss_pytorch(ref_logits, 4, score_function) - ref_map = ref_map.reshape(logits.shape) - if routing_map_format == "bitmap_u8": - ref_map = _bytemap_to_bitmap_u8(ref_map) - assert output.dtype == torch.float32 - assert not routing_map.requires_grad - torch.testing.assert_close(routing_map, eager_map) - torch.testing.assert_close(routing_map, ref_map) - torch.testing.assert_close(output, eager_output, atol=0, rtol=0) - _check_compile_reference(output, ref_output.reshape(logits.shape), torch.float32) - grad = torch.randn((*output.shape, 2), device="cuda")[..., 0] - (actual_grad,) = torch.autograd.grad(output, logits, grad) - (eager_grad,) = torch.autograd.grad(eager_output, eager_logits, grad) - (ref_grad,) = torch.autograd.grad(ref_output, ref_logits, grad.reshape(-1, 32)) - assert actual_grad.dtype == dtype - torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) - _check_compile_reference(actual_grad, ref_grad.reshape(logits.shape), dtype) - with torch.no_grad(): - torch.compiler.cudagraph_mark_step_begin() - torch.testing.assert_close(compiled(logits), scores(logits), atol=0, rtol=0) - if mode == "reduce-overhead": - assert not counters["inductor"]["cudagraph_skips"] - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) -@pytest.mark.parametrize("tensor_count", [False, True]) -def test_aux_loss_compile(dtype, mode, tensor_count): - torch._dynamo.reset() - counters.clear() - counts = torch.arange(1, 33, device="cuda", dtype=torch.int32) - total = torch.tensor(64, device="cuda", dtype=torch.int64) if tensor_count else 64 - - def loss(probs, total_num_tokens): - return fused_moe_aux_loss(probs, counts, total_num_tokens, 16, 4, 0.01) - - compiled = torch.compile(loss, fullgraph=True, mode=mode) - for new_total in (64, 128, 47): - torch.compiler.cudagraph_mark_step_begin() - if tensor_count: - total.fill_(new_total) - else: - total = new_total - probs = _compile_logits(dtype).detach().reshape(-1, 32).requires_grad_(True) - eager_probs = probs.detach().clone().requires_grad_(True) - ref_probs = probs.detach().float().requires_grad_(True) - output = compiled(probs, total) - eager_output = loss(eager_probs, total) - ref_output = aux_loss_pytorch(ref_probs, counts, new_total, 4, 16, 0.01) - assert output.shape == () and output.dtype == dtype - # The native loss reduction accumulates with FP32 atomics. - torch.testing.assert_close(output, eager_output) - _check_compile_reference(output, ref_output, dtype) - grad = output.new_tensor(0.7) - (actual_grad,) = torch.autograd.grad(output, probs, grad) - (eager_grad,) = torch.autograd.grad(eager_output, eager_probs, grad) - (ref_grad,) = torch.autograd.grad(ref_output, ref_probs, grad.float()) - torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) - _check_compile_reference(actual_grad, ref_grad, dtype) + torch.testing.assert_close(indices, torch.full_like(indices, -1 - step)) + indices.sub_(1) + histogram.add_(1) + bounds.add_(0.1) + actual_grad = torch.autograd.grad(actual[0], logits) + expected_grad = torch.autograd.grad(expected[0], eager_logits) + torch.testing.assert_close(actual_grad, expected_grad) with torch.no_grad(): torch.compiler.cudagraph_mark_step_begin() - torch.testing.assert_close(compiled(probs, total), loss(probs, total)) - if mode == "reduce-overhead": - assert not counters["inductor"]["cudagraph_skips"] + torch.testing.assert_close(compiled(logits, bounds, total), router(logits, bounds, total)) + assert not counters["inductor"]["cudagraph_skips"] -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("index_dtype", [None, torch.int16, torch.int32, torch.int64]) +@pytest.mark.parametrize( + "dtype,index_dtype", + [(torch.float32, None), (torch.float16, torch.int16), (torch.bfloat16, torch.int64)], +) def test_router_compile_dynamic(dtype, index_dtype): torch._dynamo.reset() counts = torch.arange(1, 131, device="cuda", dtype=torch.int32) @@ -358,12 +253,22 @@ def router_loss(logits): _check_compile_reference(actual_grad, eager_grad, dtype) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize( - "routing_map_format,index_dtype", - [(0, None), (1, None)] + [(0, dtype) for dtype in (torch.int16, torch.int32, torch.int64)], + "dtype,routing_map_format,index_dtype,histogram_mode", + [ + (torch.float32, 0, None, None), + (torch.float16, 1, None, None), + (torch.bfloat16, 0, torch.int16, None), + (torch.float32, 0, torch.int32, None), + (torch.bfloat16, 0, torch.int64, None), + (torch.float32, 0, None, QBHistogramMode.TWO_KERNEL), + (torch.float16, 1, None, QBHistogramMode.FUSED_ATOMIC), + (torch.bfloat16, 0, torch.int16, QBHistogramMode.TWO_KERNEL), + (torch.float32, 0, torch.int32, QBHistogramMode.FUSED_ATOMIC), + (torch.bfloat16, 0, torch.int64, QBHistogramMode.FUSED_ATOMIC), + ], ) -def test_router_compile_opcheck(dtype, routing_map_format, index_dtype): +def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogram_mode): from transformer_engine.pytorch import router logits = torch.randn(16, 130, device="cuda", dtype=dtype, requires_grad=True) @@ -372,17 +277,37 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype): if index_dtype is not None else None ) - torch.library.opcheck( - router._topk_forward, - (logits, 4, False, None, None, None, "sigmoid", None, routing_map_format, indices), - ) - torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) - counts = torch.ones(130, device="cuda", dtype=torch.int32) - count = torch.tensor(16, device="cuda", dtype=torch.int64) - for token_count in (None, count): + if histogram_mode is not None: + bias = torch.zeros(130, device="cuda") + histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.0, 1.0], device="cuda") torch.library.opcheck( - router._aux_loss_forward, (logits, counts, 16, token_count, 130, 4, 0.01) + router._qb_forward, + ( + logits, + 4, + None, + bias, + routing_map_format, + indices, + histogram, + bounds, + int(histogram_mode), + ), ) + else: + torch.library.opcheck( + router._topk_forward, + (logits, 4, False, None, None, None, "sigmoid", None, routing_map_format, indices), + ) + if index_dtype in (None, torch.int16): + torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) + counts = torch.ones(130, device="cuda", dtype=torch.int32) + count = torch.tensor(16, device="cuda", dtype=torch.int64) + for token_count in (None, count): + torch.library.opcheck( + router._aux_loss_forward, (logits, counts, 16, token_count, 130, 4, 0.01) + ) @pytest.mark.parametrize( @@ -432,196 +357,19 @@ def router(x): torch.compile(router, fullgraph=True)(logits) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("mode", ["default", "reduce-overhead"]) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -@pytest.mark.parametrize( - "routing_map_format,index_dtype", - [("bytemap", None), ("bitmap_u8", None)] - + [("bytemap", dtype) for dtype in (torch.int16, torch.int32, torch.int64)], -) -@pytest.mark.parametrize("topk", [1, 4, 16]) -def test_qb_compile(dtype, mode, histogram_mode, routing_map_format, index_dtype, topk): - torch._dynamo.reset() - counters.clear() - bias = torch.linspace(0.3, 0, 32, device="cuda") - histogram = torch.full((32, 31), 7, device="cuda", dtype=torch.int32) - bounds = torch.tensor([-1.1, 1.2], device="cuda") - indices = ( - torch.full((2, 32, topk), -1, device="cuda", dtype=index_dtype) - if index_dtype is not None - else None - ) - - def router(logits, indices, histogram, bounds): - return fused_topk_with_score_function( - logits, - topk, - False, - None, - None, - 1.2, - "sigmoid", - bias, - routing_map_format=routing_map_format, - topk_indices=indices, - qb_histogram=histogram, - qb_bin_bounds=bounds, - qb_histogram_mode=histogram_mode, - overwrite_topk_indices=False, - return_qb_histogram_delta=True, - assume_valid_qb_bin_bounds=True, - ) - - compiled = torch.compile(router, fullgraph=True, mode=mode) - accumulated = torch.zeros_like(histogram) - reference_accumulated = torch.zeros_like(histogram) - for step, new_bounds in enumerate(([-1.1, 1.2], [-0.3, 0.7], [-0.8, 1.4])): - torch.compiler.cudagraph_mark_step_begin() - bounds.copy_(bounds.new_tensor(new_bounds)) - logits = _compile_logits(dtype) - if step == 1: - logits = logits.flip(-1).detach().requires_grad_(True) - eager_logits = logits.detach().clone().requires_grad_(True) - ref_logits = logits.detach().float().requires_grad_(True) - histogram_version = histogram._version - indices_version = indices._version if indices is not None else None - probs, routing, delta = compiled(logits, indices, histogram, bounds) - eager_probs, eager_routing, eager_delta = router(eager_logits, indices, histogram, bounds) - reference = qb_topk_score_function_pytorch(ref_logits, topk, bias, bounds, 31) - ref_probs = reference["probs"] * 1.2 - ref_map = reference["routing_map"] - assert not routing.requires_grad and not delta.requires_grad - assert delta.data_ptr() != histogram.data_ptr() - assert eager_delta.data_ptr() != histogram.data_ptr() - assert histogram._version == histogram_version - torch.testing.assert_close(histogram, torch.full_like(histogram, 7 + step)) - torch.testing.assert_close(delta, eager_delta) - torch.testing.assert_close(delta, reference["histogram"]) - torch.testing.assert_close(routing, eager_routing) - if indices is not None: - assert routing.shape == indices.shape and routing.dtype == index_dtype - assert routing.data_ptr() != indices.data_ptr() - assert eager_routing.data_ptr() != indices.data_ptr() - assert indices._version == indices_version - torch.testing.assert_close(indices, torch.full_like(indices, -1 - step)) - actual_map = topk_indices_to_routing_map(routing.reshape(-1, topk), 32) - torch.testing.assert_close(actual_map.reshape(logits.shape), ref_map) - indices.sub_(1) - elif routing_map_format == "bitmap_u8": - torch.testing.assert_close(routing, _bytemap_to_bitmap_u8(ref_map)) - else: - torch.testing.assert_close(routing, ref_map) - accumulated.add_(delta) - reference_accumulated.add_(reference["histogram"]) - torch.testing.assert_close(accumulated, reference_accumulated) - torch.testing.assert_close(probs, eager_probs, atol=0, rtol=0) - _check_compile_reference(probs, ref_probs, dtype) - histogram.add_(1) - bounds.add_(0.1) - grad = torch.randn((*probs.shape, 2), dtype=dtype, device="cuda")[..., 0] - (actual_grad,) = torch.autograd.grad(probs, logits, grad) - (eager_grad,) = torch.autograd.grad(eager_probs, eager_logits, grad) - (ref_grad,) = torch.autograd.grad(ref_probs, ref_logits, grad.float()) - torch.testing.assert_close(actual_grad, eager_grad, atol=0, rtol=0) - _check_compile_reference(actual_grad, ref_grad, dtype) - with torch.no_grad(): - torch.compiler.cudagraph_mark_step_begin() - torch.testing.assert_close( - compiled(logits, indices, histogram, bounds), - router(logits, indices, histogram, bounds), - atol=0, - rtol=0, - ) - if mode == "reduce-overhead": - assert not counters["inductor"]["cudagraph_skips"] - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_compile_dynamic_accumulation(dtype, histogram_mode): - torch._dynamo.reset() - bias = torch.linspace(0.3, 0, 130, device="cuda") - bounds = torch.tensor([-1.0, 1.0], device="cuda") - histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) - reference_histogram = histogram.clone() - - def router(logits, histogram, bounds): - probs, routing, delta = fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - bias, - "bitmap_u8", - qb_histogram=histogram, - qb_bin_bounds=bounds, - qb_histogram_mode=histogram_mode, - return_qb_histogram_delta=True, - assume_valid_qb_bin_bounds=True, - ) - return probs, routing, histogram + delta - - compiled = torch.compile(router, fullgraph=True, dynamic=True) - for num_tokens in (17, 33, 64): - logits = torch.randn(num_tokens, 130, device="cuda", dtype=dtype, requires_grad=True) - input_histogram = histogram.clone() - probs, routing, histogram = compiled(logits, input_histogram, bounds) - reference = qb_topk_score_function_pytorch(logits, 4, bias, bounds, 64, reference_histogram) - torch.testing.assert_close(histogram, reference["histogram"]) - torch.testing.assert_close(routing, _bytemap_to_bitmap_u8(reference["routing_map"])) - _check_compile_reference(probs, reference["probs"], dtype) - assert histogram.data_ptr() != input_histogram.data_ptr() - torch.autograd.grad(probs.sum(), logits) - bounds.add_(0.03) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize( - "histogram_mode", [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC] -) -@pytest.mark.parametrize( - "routing_map_format,index_dtype", - [(0, None), (1, None)] + [(0, dtype) for dtype in (torch.int16, torch.int32, torch.int64)], -) -def test_qb_compile_opcheck(dtype, histogram_mode, routing_map_format, index_dtype): - from transformer_engine.pytorch import router - - logits = torch.randn(16, 130, device="cuda", dtype=dtype, requires_grad=True) - bias = torch.zeros(130, device="cuda") - histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) - bounds = torch.tensor([-1.0, 1.0], device="cuda") - indices = ( - torch.full((16, 4), -1, device="cuda", dtype=index_dtype) - if index_dtype is not None - else None - ) - torch.library.opcheck( - router._qb_forward, - ( - logits, - 4, - None, - bias, - routing_map_format, - indices, - histogram, - bounds, - int(histogram_mode), - ), - ) - - @pytest.mark.parametrize("delta", [False, True]) def test_qb_trusted_bounds_skip_value_validation(monkeypatch, delta): from transformer_engine.pytorch import router logits = torch.randn(8, 16, device="cuda", requires_grad=True) - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bias = torch.zeros(16, device="cuda") bounds = torch.tensor([-1.0, 1.0], device="cuda") + options = dict( + qb_histogram=torch.zeros(16, 32, device="cuda", dtype=torch.int32), + qb_bin_bounds=bounds, + qb_histogram_mode="two_kernel", + return_qb_histogram_delta=delta, + ) def unexpected_validation(_): raise AssertionError("trusted QB bounds must not read values on the CPU") @@ -636,28 +384,14 @@ def unexpected_validation(_): None, None, "sigmoid", - torch.zeros(16, device="cuda"), - qb_histogram=histogram, - qb_bin_bounds=bounds, - qb_histogram_mode="two_kernel", - return_qb_histogram_delta=delta, + bias, + **options, assume_valid_qb_bin_bounds=True, ) bounds.zero_() with pytest.raises(ValueError, match="finite with lower < upper"): fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - torch.zeros(16, device="cuda"), - qb_histogram=histogram, - qb_bin_bounds=bounds, - qb_histogram_mode="two_kernel", - return_qb_histogram_delta=delta, + logits, 4, False, None, None, None, "sigmoid", bias, **options ) @@ -699,22 +433,6 @@ def router(logits, bounds): router(logits, bounds) -def test_qb_histogram_delta_requires_qb(): - logits = torch.randn(8, 16, device="cuda") - with pytest.raises(ValueError, match="return_qb_histogram_delta requires"): - fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - None, - return_qb_histogram_delta=True, - ) - - # Pytorch-based group topk def group_limited_topk( scores: torch.Tensor, @@ -918,6 +636,7 @@ def run_comparison( enable_bias, topk_output_mode="sparse", topk_index_dtype=torch.int16, + execution="eager", ): if topk >= num_experts: pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") @@ -946,9 +665,12 @@ def run_comparison( * 1e-4 ) logits = logits.view(num_tokens, num_experts) + if dtype != torch.float32: + logits = torch.linspace(-2, 2, num_experts, device="cuda").repeat(num_tokens, 1) + logits = logits.to(dtype) logits.requires_grad = True if enable_bias and score_function in ("sigmoid", "sqrtsoftplus"): - expert_bias = torch.arange(num_experts, device="cuda", dtype=dtype) * 0.1 + expert_bias = torch.arange(num_experts, device="cuda", dtype=torch.float32) * 0.1 expert_bias = torch.flip(expert_bias, dims=[0]) expert_bias.requires_grad = True else: @@ -963,6 +685,8 @@ def run_comparison( else: expert_bias_clone = None + logits = logits.detach().float().requires_grad_(True) + # Run the original implementation # We do not support the capacity factor case probs, routing_map = topk_score_function_pytorch( @@ -978,10 +702,11 @@ def run_comparison( topk_indices = None if topk_output_mode == "dense": - topk_indices = torch.empty((num_tokens, topk), device="cuda", dtype=topk_index_dtype) + topk_indices = torch.full((num_tokens, topk), -1, device="cuda", dtype=topk_index_dtype) # Run the fused implementation - probs_fused, routing_output_fused = fused_topk_with_score_function( + router = _router_function(fused_topk_with_score_function, execution) + probs_fused, routing_output_fused = router( logits=logits_clone, topk=topk, use_pre_softmax=use_pre_softmax, @@ -991,38 +716,56 @@ def run_comparison( score_function=score_function, expert_bias=expert_bias_clone, topk_indices=topk_indices, + overwrite_topk_indices=execution == "eager", ) + assert probs_fused.dtype == dtype and not routing_output_fused.requires_grad if topk_output_mode == "dense": - assert routing_output_fused.data_ptr() == topk_indices.data_ptr() + assert (routing_output_fused.data_ptr() == topk_indices.data_ptr()) == ( + execution == "eager" + ) + if execution != "eager": + torch.testing.assert_close(topk_indices, torch.full_like(topk_indices, -1)) + topk_indices.fill_(-2) assert routing_output_fused.dtype == topk_index_dtype routing_map_fused = topk_indices_to_routing_map(routing_output_fused, num_experts) else: routing_map_fused = routing_output_fused atol, rtol = _get_tolerances(dtype, num_experts) - torch.testing.assert_close(probs, probs_fused, atol=atol, rtol=rtol) + torch.testing.assert_close(probs, probs_fused.float(), atol=atol, rtol=rtol) torch.testing.assert_close(routing_map, routing_map_fused) - # Fake the loss - loss = torch.sum(probs) - loss_fused = torch.sum(probs_fused) - - # Backward the loss - loss.backward() - loss_fused.backward() - - # Check the gradient - torch.testing.assert_close(logits.grad, logits_clone.grad, atol=atol, rtol=rtol) + grad = torch.randn((*probs_fused.shape, 2), device="cuda", dtype=dtype)[..., 0] + probs.backward(grad.float()) + probs_fused.backward(grad) + torch.testing.assert_close(logits.grad, logits_clone.grad.float(), atol=atol, rtol=rtol) + + +_SIGMOID_TOPK_CASES = [ + ("eager", *case) + for case in product( + [torch.float32], + [2048, 7168, 8992], + [1024, 128, 32], + [4, 8, 16, 32], + [None, 4], + [None, 1.2], + [True, False], + [None, torch.int16, torch.int32, torch.int64], + ) +] + [ + ("compile", torch.float32, 64, 32, 4, None, None, False, None), + ("compile", torch.float16, 64, 32, 4, None, 1.2, True, torch.int16), + ("compile", torch.bfloat16, 64, 32, 4, 4, 1.2, True, torch.int32), + ("compile", torch.float32, 64, 130, 16, None, None, True, torch.int64), + ("compile", torch.bfloat16, 64, 32, 1, None, None, False, None), +] -@pytest.mark.parametrize("dtype", [torch.float32]) -@pytest.mark.parametrize("num_tokens", [2048, 7168, 8992]) -@pytest.mark.parametrize("num_experts", [1024, 128, 32]) -@pytest.mark.parametrize("topk", [4, 8, 16, 32]) -@pytest.mark.parametrize("group_topk", [None, 4]) -@pytest.mark.parametrize("scaling_factor", [None, 1.2]) -@pytest.mark.parametrize("enable_bias", [True, False]) -@pytest.mark.parametrize("topk_index_dtype", [None, torch.int16, torch.int32, torch.int64]) +@pytest.mark.parametrize( + "execution,dtype,num_tokens,num_experts,topk,group_topk,scaling_factor,enable_bias,topk_index_dtype", + _SIGMOID_TOPK_CASES, +) def test_topk_sigmoid( dtype, num_tokens, @@ -1032,6 +775,7 @@ def test_topk_sigmoid( scaling_factor, enable_bias, topk_index_dtype, + execution, ): num_groups = 8 if group_topk else None run_comparison( @@ -1047,17 +791,14 @@ def test_topk_sigmoid( enable_bias=enable_bias, topk_output_mode="dense" if topk_index_dtype is not None else "sparse", topk_index_dtype=topk_index_dtype or torch.int16, + execution=execution, ) -@pytest.mark.parametrize("dtype", [torch.float32]) -@pytest.mark.parametrize("num_tokens", [2048, 7168, 8992]) -@pytest.mark.parametrize("num_experts", [1024, 128, 32]) -@pytest.mark.parametrize("topk", [4, 8, 16, 32]) -@pytest.mark.parametrize("group_topk", [None, 4]) -@pytest.mark.parametrize("scaling_factor", [None, 1.2]) -@pytest.mark.parametrize("enable_bias", [True, False]) -@pytest.mark.parametrize("topk_index_dtype", [None, torch.int16, torch.int32, torch.int64]) +@pytest.mark.parametrize( + "execution,dtype,num_tokens,num_experts,topk,group_topk,scaling_factor,enable_bias,topk_index_dtype", + _SIGMOID_TOPK_CASES, +) def test_topk_sqrtsoftplus( dtype, num_tokens, @@ -1067,6 +808,7 @@ def test_topk_sqrtsoftplus( scaling_factor, enable_bias, topk_index_dtype, + execution, ): num_groups = 8 if group_topk else None run_comparison( @@ -1082,28 +824,57 @@ def test_topk_sqrtsoftplus( enable_bias=enable_bias, topk_output_mode="dense" if topk_index_dtype is not None else "sparse", topk_index_dtype=topk_index_dtype or torch.int16, + execution=execution, ) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -@pytest.mark.parametrize("topk", [8, 16]) @pytest.mark.parametrize( - "routing_output_mode", - ["bytemap", "bitmap_u8", "dense_int16", "dense_int32", "dense_int64"], + "execution,dtype,histogram_mode,topk,routing_output_mode,overwrite,delta", + [ + ("eager", torch.float32, *case, overwrite, False) + for case in product( + ["two_kernel", "fused_atomic"], + [8, 16], + ["bytemap", "bitmap_u8", "dense_int16", "dense_int32", "dense_int64"], + ) + for overwrite in ([False, True] if case[2].startswith("dense") else [True]) + ] + + [ + (execution, dtype, histogram_mode, topk, layout, False, True) + for execution in ["eager", "compile"] + for dtype, histogram_mode, topk, layout in [ + (torch.float32, "two_kernel", 1, "bytemap"), + (torch.float16, "fused_atomic", 4, "bitmap_u8"), + (torch.bfloat16, "two_kernel", 16, "dense_int16"), + (torch.float32, "fused_atomic", 16, "dense_int32"), + (torch.bfloat16, "fused_atomic", 4, "dense_int64"), + (torch.float16, "two_kernel", 4, "bytemap"), + ] + ], ) -@pytest.mark.parametrize("overwrite", [False, True]) -def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, overwrite): - num_tokens = 257 - num_experts = 896 - num_bins = 1000 +def test_qb_topk_histogram( + histogram_mode, + topk, + routing_output_mode, + overwrite, + execution, + dtype, + delta, +): + num_tokens, num_experts, num_bins = (64, 32, 31) if delta else (257, 896, 1000) logits = torch.randn( num_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True ) + if delta: + logits = ( + _compile_logits(dtype).reshape(num_tokens, num_experts).detach().requires_grad_(True) + ) + reference_logits = logits.detach().float().requires_grad_(True) expert_bias = torch.linspace(-0.2, 0.2, num_experts, device="cuda", dtype=torch.float32) bin_bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) reference_histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) reference = qb_topk_score_function_pytorch( - logits, + reference_logits, topk, expert_bias, bin_bounds, @@ -1112,7 +883,8 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, overwrite) ) fused_logits = logits.detach().clone().requires_grad_(True) - fused_histogram = torch.zeros_like(reference_histogram) + fused_histogram = torch.full_like(reference_histogram, 7 if delta else 0) + histogram_version = fused_histogram._version dense_dtype = { "dense_int16": torch.int16, "dense_int32": torch.int32, @@ -1128,7 +900,8 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, overwrite) if routing_output_mode == "bitmap_u8" else RoutingMapFormat.BYTEMAP ) - fused_probs, fused_routing_output = fused_topk_with_score_function( + router = _router_function(fused_topk_with_score_function, execution) + outputs = router( logits=fused_logits, topk=topk, use_pre_softmax=False, @@ -1143,12 +916,30 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, overwrite) qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, overwrite_topk_indices=overwrite, + return_qb_histogram_delta=delta, + assume_valid_qb_bin_bounds=delta, ) - torch.testing.assert_close(fused_probs, reference["probs"]) + fused_probs, fused_routing_output = outputs[:2] + if delta: + _check_compile_reference(fused_probs, reference["probs"], dtype) + else: + torch.testing.assert_close(fused_probs, reference["probs"]) + assert not fused_routing_output.requires_grad + if delta: + assert not outputs[2].requires_grad + assert outputs[2].data_ptr() != fused_histogram.data_ptr() + assert fused_histogram._version == histogram_version + torch.testing.assert_close(fused_histogram, torch.full_like(fused_histogram, 7)) + torch.testing.assert_close(outputs[2], reference["histogram"]) + fused_histogram.add_(1) + bin_bounds.add_(0.1) + else: + torch.testing.assert_close(fused_histogram, reference["histogram"]) if dense_dtype is not None: assert (fused_routing_output.data_ptr() == topk_indices.data_ptr()) == overwrite if not overwrite: torch.testing.assert_close(topk_indices, torch.full_like(topk_indices, -1)) + topk_indices.fill_(-2) fused_routing_map = topk_indices_to_routing_map(fused_routing_output, num_experts) torch.testing.assert_close(fused_routing_map, reference["routing_map"]) elif routing_output_mode == "bitmap_u8": @@ -1158,49 +949,80 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, overwrite) ) else: torch.testing.assert_close(fused_routing_output, reference["routing_map"]) - torch.testing.assert_close(fused_histogram, reference["histogram"]) - grad = torch.randn_like(fused_probs) - reference["probs"].backward(grad) + grad = torch.randn((*fused_probs.shape, 2), device="cuda", dtype=dtype)[..., 0] + reference["probs"].backward(grad.float()) fused_probs.backward(grad) - torch.testing.assert_close(fused_logits.grad, logits.grad, atol=1e-5, rtol=1e-5) + if delta: + _check_compile_reference(fused_logits.grad, reference_logits.grad, dtype) + else: + torch.testing.assert_close(fused_logits.grad, reference_logits.grad, atol=1e-5, rtol=1e-5) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_histogram_accumulates_microbatches(histogram_mode): - num_experts = 64 - topk = 8 - num_bins = 1000 - expert_bias = torch.linspace(-0.1, 0.1, num_experts, device="cuda", dtype=torch.float32) - bin_bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) - reference_histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) - fused_histogram = torch.zeros_like(reference_histogram) +@pytest.mark.parametrize( + "execution,dtype,histogram_mode", + [ + ("eager", torch.float32, "two_kernel"), + ("eager", torch.float32, "fused_atomic"), + ("compile", torch.float32, "two_kernel"), + ("compile", torch.bfloat16, "fused_atomic"), + ], +) +def test_qb_histogram_accumulates_microbatches(execution, dtype, histogram_mode): + functional = execution == "compile" + num_experts, num_bins = (130, 64) if functional else (64, 1000) + expert_bias = torch.linspace(-0.1, 0.1, num_experts, device="cuda") + bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) + histogram = torch.full( + (num_experts, num_bins), 7 if functional else 0, device="cuda", dtype=torch.int32 + ) + reference_histogram = histogram.clone() - for num_tokens in (127, 193): - logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32) - qb_topk_score_function_pytorch( + def route(logits, histogram, bounds): + outputs = fused_topk_with_score_function( logits, - topk, + 8, + False, + None, + None, + None, + "sigmoid", expert_bias, - bin_bounds, - num_bins, - reference_histogram, - ) - fused_topk_with_score_function( - logits=logits, - topk=topk, - use_pre_softmax=False, - num_groups=None, - group_topk=None, - scaling_factor=None, - score_function="sigmoid", - expert_bias=expert_bias, - qb_histogram=fused_histogram, - qb_bin_bounds=bin_bounds, + routing_map_format="bitmap_u8" if functional else "bytemap", + qb_histogram=histogram, + qb_bin_bounds=bounds, qb_histogram_mode=histogram_mode, + return_qb_histogram_delta=functional, + assume_valid_qb_bin_bounds=functional, ) - - torch.testing.assert_close(fused_histogram, reference_histogram) + return outputs[0], outputs[1], histogram + outputs[2] if functional else histogram + + if functional: + torch._dynamo.reset() + route = torch.compile(route, fullgraph=True, dynamic=True) + for num_tokens in (17, 33, 64) if functional else (127, 193): + logits = torch.randn( + num_tokens, num_experts, device="cuda", dtype=dtype, requires_grad=True + ) + reference = qb_topk_score_function_pytorch( + logits, 8, expert_bias, bounds, num_bins, reference_histogram + ) + input_histogram = histogram + probs, routing, histogram = route(logits, histogram, bounds) + torch.testing.assert_close(histogram, reference["histogram"]) + _check_compile_reference(probs, reference["probs"], dtype) + expected_map = ( + _bytemap_to_bitmap_u8(reference["routing_map"]) + if functional + else reference["routing_map"] + ) + torch.testing.assert_close(routing, expected_map) + if functional: + assert histogram.data_ptr() != input_histogram.data_ptr() + actual_grad = torch.autograd.grad(probs.float().square().sum(), logits)[0] + ref_grad = torch.autograd.grad(reference["probs"].float().square().sum(), logits)[0] + _check_compile_reference(actual_grad, ref_grad, dtype) + bounds.add_(0.03) def test_qb_topk_argument_validation(): @@ -1235,6 +1057,19 @@ def test_qb_topk_argument_validation(): qb_histogram_mode="two_kernel", ) + with pytest.raises(ValueError, match="return_qb_histogram_delta requires"): + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + None, + return_qb_histogram_delta=True, + ) + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least two CUDA devices") @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @@ -1626,14 +1461,29 @@ def test_qb_topk_plus_one_tie_and_bin_clamping(histogram_mode): assert routing_map[0, :3].tolist() == [True, True, False] -@pytest.mark.parametrize("dtype", [torch.float32]) -@pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) -@pytest.mark.parametrize("num_experts", [1024, 128, 32]) -@pytest.mark.parametrize("topk", [4, 8, 16, 32]) -@pytest.mark.parametrize("use_pre_softmax", [True, False]) -@pytest.mark.parametrize("group_topk", [None, 4]) -@pytest.mark.parametrize("scaling_factor", [None, 1.2]) -@pytest.mark.parametrize("topk_index_dtype", [None, torch.int16, torch.int32, torch.int64]) +@pytest.mark.parametrize( + "execution,dtype,num_tokens,num_experts,topk,use_pre_softmax,group_topk,scaling_factor,topk_index_dtype", + [ + ("eager", *case) + for case in product( + [torch.float32], + [2048, 7168, 14234], + [1024, 128, 32], + [4, 8, 16, 32], + [True, False], + [None, 4], + [None, 1.2], + [None, torch.int16, torch.int32, torch.int64], + ) + ] + + [ + ("compile", torch.float32, 64, 32, 4, True, None, None, None), + ("compile", torch.float16, 64, 32, 4, False, None, 1.2, torch.int16), + ("compile", torch.bfloat16, 64, 32, 4, False, 4, 1.2, torch.int32), + ("compile", torch.float32, 64, 130, 16, True, None, None, torch.int64), + ("compile", torch.bfloat16, 64, 32, 4, True, None, 1.2, None), + ], +) def test_topk_softmax( dtype, num_tokens, @@ -1643,6 +1493,7 @@ def test_topk_softmax( group_topk, scaling_factor, topk_index_dtype, + execution, ): num_groups = 8 if group_topk else None run_comparison( @@ -1658,11 +1509,14 @@ def test_topk_softmax( enable_bias=False, topk_output_mode="dense" if topk_index_dtype is not None else "sparse", topk_index_dtype=topk_index_dtype or torch.int16, + execution=execution, ) -@pytest.mark.parametrize("topk_index_dtype", [None, torch.int16]) -def test_topk_preserves_leading_dims(topk_index_dtype): +@pytest.mark.parametrize( + "execution,topk_index_dtype", product(["eager", "compile"], [None, torch.int16]) +) +def test_topk_preserves_leading_dims(topk_index_dtype, execution): num_tokens = 128 num_experts = 32 topk = 4 @@ -1671,7 +1525,8 @@ def test_topk_preserves_leading_dims(topk_index_dtype): if topk_index_dtype is not None: topk_indices = torch.empty(num_tokens, 2, topk, device="cuda", dtype=topk_index_dtype) - probs, routing_output = fused_topk_with_score_function( + router = _router_function(fused_topk_with_score_function, execution) + probs, routing_output = router( logits=logits, topk=topk, use_pre_softmax=False, @@ -1681,6 +1536,7 @@ def test_topk_preserves_leading_dims(topk_index_dtype): score_function="softmax", expert_bias=None, topk_indices=topk_indices, + overwrite_topk_indices=execution == "eager", ) assert probs.shape == logits.shape @@ -1688,12 +1544,37 @@ def test_topk_preserves_leading_dims(topk_index_dtype): assert routing_output.shape == expected_routing_shape -@pytest.mark.parametrize("dtype", [torch.float32]) -@pytest.mark.parametrize("num_tokens", [2048, 7168]) -@pytest.mark.parametrize("num_experts", [1024, 256, 128, 32]) -@pytest.mark.parametrize("topk", [1, 4, 8, 16, 32]) -@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) -def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function): +@pytest.mark.parametrize( + "execution,dtype,num_tokens,num_experts,topk,score_function,routing_map_format", + [ + ("eager", *case, "bytemap") + for case in product( + [torch.float32], + [2048, 7168], + [1024, 256, 128, 32], + [1, 4, 8, 16, 32], + ["softmax", "sigmoid", "sqrtsoftplus"], + ) + ] + + [ + ("compile", dtype, 64, 32, 4, score, "bitmap_u8" if i % 2 else "bytemap") + for i, (dtype, score) in enumerate( + product( + [torch.float32, torch.float16, torch.bfloat16], + ["softmax", "sigmoid", "sqrtsoftplus"], + ) + ) + ], +) +def test_fused_scores_for_aux_loss( + dtype, + num_tokens, + num_experts, + topk, + score_function, + execution, + routing_map_format, +): if topk >= num_experts: pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") if score_function in ("sigmoid", "sqrtsoftplus"): @@ -1714,41 +1595,62 @@ def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_f * 1e-4 ) logits = logits.view(num_tokens, num_experts) + if dtype != torch.float32: + logits = torch.linspace(-2, 2, num_experts, device="cuda").repeat(num_tokens, 1).to(dtype) logits.requires_grad = True logits_clone = deepcopy(logits) logits_clone.requires_grad = True + logits = logits.detach().float().requires_grad_(True) routing_map, scores = compute_scores_for_aux_loss_pytorch( logits=logits, topk=topk, score_function=score_function, ) - routing_map_fused, scores_fused = fused_compute_score_for_moe_aux_loss( + scores_fn = _router_function(fused_compute_score_for_moe_aux_loss, execution) + routing_map_fused, scores_fused = scores_fn( logits=logits_clone, topk=topk, score_function=score_function, + routing_map_format=routing_map_format, ) - atol, rtol = _get_tolerances(dtype, num_experts) + assert scores_fused.dtype == torch.float32 + assert not routing_map_fused.requires_grad + if routing_map_format == "bitmap_u8": + routing_map = _bytemap_to_bitmap_u8(routing_map) + atol, rtol = _get_tolerances(torch.float32, num_experts) torch.testing.assert_close(scores, scores_fused, atol=atol, rtol=rtol) torch.testing.assert_close(routing_map, routing_map_fused) - loss = torch.sum(scores) - loss.backward() - loss_fused = torch.sum(scores_fused) - loss_fused.backward() - - torch.testing.assert_close(logits.grad, logits_clone.grad, atol=atol, rtol=rtol) + grad = torch.randn((*scores_fused.shape, 2), device="cuda")[..., 0] + scores.backward(grad) + scores_fused.backward(grad) + assert logits_clone.grad.dtype == dtype + atol, rtol = _get_tolerances(dtype, num_experts) + torch.testing.assert_close(logits.grad, logits_clone.grad.float(), atol=atol, rtol=rtol) -@pytest.mark.parametrize("dtype", [torch.float32]) -@pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) -@pytest.mark.parametrize("num_experts", [1024, 256, 128, 32]) -@pytest.mark.parametrize("topk", [4, 32]) -@pytest.mark.parametrize("expert_multiplier", [1, 2]) -def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multiplier): +@pytest.mark.parametrize( + "execution,dtype,num_tokens,num_experts,topk,expert_multiplier,tensor_count", + [ + ("eager", *case, False) + for case in product( + [torch.float32], [2048, 7168, 14234], [1024, 256, 128, 32], [4, 32], [1, 2] + ) + ] + + [ + ("compile", dtype, 64, 32, 4, 2 if tensor_count else 1, tensor_count) + for dtype, tensor_count in product( + [torch.float32, torch.float16, torch.bfloat16], [False, True] + ) + ], +) +def test_fused_moe_aux_loss( + dtype, num_tokens, num_experts, topk, expert_multiplier, execution, tensor_count +): if topk >= num_experts: pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") # Sequence aux loss batches independent sequences along the expert dimension. @@ -1766,6 +1668,7 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl probs_clone = deepcopy(probs) probs_clone.requires_grad = True + probs = probs.detach().float().requires_grad_(True) aux_loss = aux_loss_pytorch( probs=probs, tokens_per_expert=tokens_per_expert, @@ -1775,23 +1678,28 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl moe_aux_loss_coeff=coeff, ) - aux_loss_fused = fused_moe_aux_loss( + total = ( + torch.tensor(num_tokens, device="cuda", dtype=torch.int64) if tensor_count else num_tokens + ) + loss_fn = _router_function(fused_moe_aux_loss, execution) + aux_loss_fused = loss_fn( probs=probs_clone, tokens_per_expert=tokens_per_expert, - total_num_tokens=num_tokens, + total_num_tokens=total, num_experts=num_experts, topk=topk, coeff=coeff, ) + assert aux_loss_fused.shape == () and aux_loss_fused.dtype == dtype atol, rtol = _get_tolerances(dtype, num_cols) - torch.testing.assert_close(aux_loss, aux_loss_fused, atol=atol, rtol=rtol) + torch.testing.assert_close(aux_loss, aux_loss_fused.float(), atol=atol, rtol=rtol) # Backward aux_loss.backward() aux_loss_fused.backward() - torch.testing.assert_close(probs.grad, probs_clone.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(probs.grad, probs_clone.grad.float(), atol=atol, rtol=rtol) def test_fused_moe_aux_loss_cuda_graph_capture(): @@ -1875,13 +1783,21 @@ def _bytemap_to_bitmap_u8(bytemap: torch.Tensor) -> torch.Tensor: return torch.from_numpy(np.packbits(flat, axis=-1, bitorder="little")).to(bytemap.device) -@pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize( - "num_tokens,num_experts,topk", - [(128, 32, 4), (256, 128, 8), (256, 130, 8), (128, 1024, 16)], + "execution,dtype,num_tokens,num_experts,topk,score_function", + [ + ("eager", torch.float32, *shape, score) + for shape, score in product( + [(128, 32, 4), (256, 128, 8), (256, 130, 8), (128, 1024, 16)], + ["softmax", "sigmoid", "sqrtsoftplus"], + ) + ] + + [ + ("compile", torch.float32, 64, 130, 4, score) + for score in ["softmax", "sigmoid", "sqrtsoftplus"] + ], ) -@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) -def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): +def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function, execution): """fused_topk_with_score_function should produce identical probs and an LSB-packed bitmap routing_map when routing_map_format=BITMAP_U8, and backward gradients should match the bytemap path exactly.""" @@ -1908,7 +1824,8 @@ def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_func logits_byte = logits.detach().clone().requires_grad_(True) logits_bit = logits.detach().clone().requires_grad_(True) - probs_byte, routing_map_byte = fused_topk_with_score_function( + router = _router_function(fused_topk_with_score_function, execution) + probs_byte, routing_map_byte = router( logits=logits_byte, topk=topk, use_pre_softmax=False, @@ -1919,7 +1836,7 @@ def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_func expert_bias=None, routing_map_format=RoutingMapFormat.BYTEMAP, ) - probs_bit, routing_map_bit = fused_topk_with_score_function( + probs_bit, routing_map_bit = router( logits=logits_bit, topk=topk, use_pre_softmax=False, @@ -1944,19 +1861,29 @@ def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_func packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) torch.testing.assert_close(routing_map_bit, packed_expected, atol=0, rtol=0) - # Backward parity: grad of probs.sum() must be bit-identical across formats. - probs_byte.sum().backward() - probs_bit.sum().backward() + # Backward must be bit-identical across formats. + grad = torch.randn_like(probs_byte) + probs_byte.backward(grad) + probs_bit.backward(grad) torch.testing.assert_close(logits_byte.grad, logits_bit.grad, atol=0.0, rtol=0.0) -@pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize( - "num_tokens,num_experts,topk", - [(128, 32, 4), (256, 128, 8), (256, 130, 8)], + "execution,dtype,num_tokens,num_experts,topk,score_function", + [ + ("eager", torch.float32, *shape, score) + for shape, score in product( + [(128, 32, 4), (256, 128, 8), (256, 130, 8)], ["softmax", "sigmoid", "sqrtsoftplus"] + ) + ] + + [ + ("compile", torch.float32, 64, 130, 4, score) + for score in ["softmax", "sigmoid", "sqrtsoftplus"] + ], ) -@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) -def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): +def test_score_for_aux_loss_bitmap_vs_bytemap( + dtype, num_tokens, num_experts, topk, score_function, execution +): """fused_compute_score_for_moe_aux_loss: bitmap routing_map must equal LSB-packed bytemap; scores must be bit-identical across formats.""" if topk >= num_experts: @@ -1968,13 +1895,14 @@ def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, to logits_byte = logits.detach().clone().requires_grad_(True) logits_bit = logits.detach().clone().requires_grad_(True) - routing_map_byte, scores_byte = fused_compute_score_for_moe_aux_loss( + router = _router_function(fused_compute_score_for_moe_aux_loss, execution) + routing_map_byte, scores_byte = router( logits=logits_byte, topk=topk, score_function=score_function, routing_map_format="bytemap", ) - routing_map_bit, scores_bit = fused_compute_score_for_moe_aux_loss( + routing_map_bit, scores_bit = router( logits=logits_bit, topk=topk, score_function=score_function, @@ -1991,8 +1919,9 @@ def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, to torch.testing.assert_close(routing_map_bit, packed_expected, atol=0, rtol=0) # Backward parity through scores. - scores_byte.sum().backward() - scores_bit.sum().backward() + grad = torch.randn_like(scores_byte) + scores_byte.backward(grad) + scores_bit.backward(grad) torch.testing.assert_close(logits_byte.grad, logits_bit.grad, atol=0.0, rtol=0.0) From e0aa561b35aad1c1a74ad3b5a728884415b27e49 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 15:01:16 +0200 Subject: [PATCH 05/14] Compile router buffer mutations through autograd.Function Preserve dense-index output buffers and QB histogram accumulation under fullgraph compilation. Keep direct native eager calls and use mutable custom ops inside the existing autograd Functions for compiled training, with a direct operator path for compiled inference. Remove the functional-output flags, preserve bounds validation, and cover buffer aliasing, accumulation, saved-index lifetime and CUDA graph replay. Validation: router 3390 passed/446 skipped; pre-commit, license and scoped pylint passed. Full Python lint retains unrelated existing C0411. Native pointer probes confirm auxiliary buffer reuse; six eager measurements show no material regression against the previous PR HEAD. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 256 +++++++-------- transformer_engine/pytorch/router.py | 459 +++++++++++++++++---------- 2 files changed, 420 insertions(+), 295 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 7ea10af2f8d..f6776af3ebd 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -158,7 +158,7 @@ def test_router_compile_replay(histogram_mode, index_dtype): else None ) - def router(logits, bounds, total): + def router(logits, bounds, total, indices, histogram): outputs = fused_topk_with_score_function( logits, 4, @@ -170,17 +170,20 @@ def router(logits, bounds, total): bias, routing_map_format="bitmap_u8" if indices is None else "bytemap", topk_indices=indices, - overwrite_topk_indices=False, qb_histogram=histogram if histogram_mode else None, qb_bin_bounds=bounds if histogram_mode else None, qb_histogram_mode=histogram_mode, - return_qb_histogram_delta=histogram_mode is not None, assume_valid_qb_bin_bounds=histogram_mode is not None, ) _, scores = fused_compute_score_for_moe_aux_loss(logits, 4, "sigmoid") loss = fused_moe_aux_loss(scores.reshape(-1, 32), counts, total, 32, 4, 0.01) return (outputs[0].float().square().sum() + loss, *outputs[1:]) + reference_histogram = histogram.clone() + reference_indices = torch.empty_like(indices) if indices is not None else None + torch._dynamo.mark_static_address(histogram) + if indices is not None: + torch._dynamo.mark_static_address(indices) compiled = torch.compile(router, fullgraph=True, mode="reduce-overhead") for step, new_bounds in enumerate(([-1.1, 1.2], [-0.3, 0.7], [-0.8, 1.4])): torch.compiler.cudagraph_mark_step_begin() @@ -190,21 +193,25 @@ def router(logits, bounds, total): if step == 1: logits = logits.flip(-1).detach().requires_grad_(True) eager_logits = logits.detach().clone().requires_grad_(True) - actual = compiled(logits, bounds, total) - expected = router(eager_logits, bounds, total) + actual = compiled(logits, bounds, total, indices, histogram) + expected = router(eager_logits, bounds, total, reference_indices, reference_histogram) torch.testing.assert_close(actual, expected) - torch.testing.assert_close(histogram, torch.full_like(histogram, 7 + step)) + torch.testing.assert_close(histogram, reference_histogram) if indices is not None: - torch.testing.assert_close(indices, torch.full_like(indices, -1 - step)) - indices.sub_(1) + assert actual[1].data_ptr() == indices.data_ptr() + torch.testing.assert_close(indices, reference_indices) histogram.add_(1) + reference_histogram.add_(1) bounds.add_(0.1) actual_grad = torch.autograd.grad(actual[0], logits) expected_grad = torch.autograd.grad(expected[0], eager_logits) torch.testing.assert_close(actual_grad, expected_grad) with torch.no_grad(): torch.compiler.cudagraph_mark_step_begin() - torch.testing.assert_close(compiled(logits, bounds, total), router(logits, bounds, total)) + actual = compiled(logits, bounds, total, indices, histogram) + expected = router(logits, bounds, total, reference_indices, reference_histogram) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(histogram, reference_histogram) assert not counters["inductor"]["cudagraph_skips"] @@ -233,7 +240,6 @@ def router_loss(logits): None, "bytemap" if indices is not None else "bitmap_u8", topk_indices=indices, - overwrite_topk_indices=False, ) _, scores = fused_compute_score_for_moe_aux_loss(logits, 4, "sigmoid", 1) aux_loss = fused_moe_aux_loss(scores, counts, logits.shape[0], 130, 4, 0.01) @@ -277,19 +283,20 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogra if index_dtype is not None else None ) + forward_logits = logits.detach() if histogram_mode is not None: bias = torch.zeros(130, device="cuda") histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) bounds = torch.tensor([-1.0, 1.0], device="cuda") torch.library.opcheck( - router._qb_forward, + router._qb_indices_forward if indices is not None else router._qb_map_forward, ( - logits, + forward_logits, 4, None, bias, routing_map_format, - indices, + *((indices,) if indices is not None else ()), histogram, bounds, int(histogram_mode), @@ -297,8 +304,19 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogra ) else: torch.library.opcheck( - router._topk_forward, - (logits, 4, False, None, None, None, "sigmoid", None, routing_map_format, indices), + router._topk_indices_forward if indices is not None else router._topk_map_forward, + ( + forward_logits, + 4, + False, + None, + None, + None, + "sigmoid", + None, + routing_map_format, + *((indices,) if indices is not None else ()), + ), ) if index_dtype in (None, torch.int16): torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) @@ -310,18 +328,8 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogra ) -@pytest.mark.parametrize( - "qb,overwrite,delta,assume_valid", - [ - (False, True, False, False), - (True, True, False, False), - (True, False, False, False), - (True, False, False, True), - (True, False, True, False), - (True, True, True, True), - ], -) -def test_router_compile_eager_fallback(qb, overwrite, delta, assume_valid): +@pytest.mark.parametrize("index_dtype", [None, torch.int32]) +def test_router_compile_eager_fallback(index_dtype): torch._dynamo.reset() logits = torch.randn(32, 32, device="cuda", requires_grad=True) bias = torch.zeros(32, device="cuda") @@ -339,13 +347,10 @@ def router(x): None, "sigmoid", bias, - topk_indices=indices if not qb or (delta and assume_valid) else None, - qb_histogram=histogram if qb else None, - qb_bin_bounds=bounds if qb else None, - qb_histogram_mode="two_kernel" if qb else None, - **({} if overwrite else {"overwrite_topk_indices": False}), - return_qb_histogram_delta=delta, - assume_valid_qb_bin_bounds=assume_valid, + topk_indices=indices if index_dtype is not None else None, + qb_histogram=histogram, + qb_bin_bounds=bounds, + qb_histogram_mode="two_kernel", ) with pytest.warns(UserWarning, match="Falling back to eager execution"): @@ -357,8 +362,7 @@ def router(x): torch.compile(router, fullgraph=True)(logits) -@pytest.mark.parametrize("delta", [False, True]) -def test_qb_trusted_bounds_skip_value_validation(monkeypatch, delta): +def test_qb_trusted_bounds_skip_value_validation(monkeypatch): from transformer_engine.pytorch import router logits = torch.randn(8, 16, device="cuda", requires_grad=True) @@ -368,7 +372,6 @@ def test_qb_trusted_bounds_skip_value_validation(monkeypatch, delta): qb_histogram=torch.zeros(16, 32, device="cuda", dtype=torch.int32), qb_bin_bounds=bounds, qb_histogram_mode="two_kernel", - return_qb_histogram_delta=delta, ) def unexpected_validation(_): @@ -423,7 +426,6 @@ def router(logits, bounds): qb_histogram=histogram, qb_bin_bounds=bounds, qb_histogram_mode="two_kernel", - return_qb_histogram_delta=True, assume_valid_qb_bin_bounds=True, ) @@ -716,16 +718,10 @@ def run_comparison( score_function=score_function, expert_bias=expert_bias_clone, topk_indices=topk_indices, - overwrite_topk_indices=execution == "eager", ) assert probs_fused.dtype == dtype and not routing_output_fused.requires_grad if topk_output_mode == "dense": - assert (routing_output_fused.data_ptr() == topk_indices.data_ptr()) == ( - execution == "eager" - ) - if execution != "eager": - torch.testing.assert_close(topk_indices, torch.full_like(topk_indices, -1)) - topk_indices.fill_(-2) + assert routing_output_fused.data_ptr() == topk_indices.data_ptr() assert routing_output_fused.dtype == topk_index_dtype routing_map_fused = topk_indices_to_routing_map(routing_output_fused, num_experts) else: @@ -829,18 +825,17 @@ def test_topk_sqrtsoftplus( @pytest.mark.parametrize( - "execution,dtype,histogram_mode,topk,routing_output_mode,overwrite,delta", + "execution,dtype,histogram_mode,topk,routing_output_mode,shape", [ - ("eager", torch.float32, *case, overwrite, False) + ("eager", torch.float32, *case, (257, 896, 1000)) for case in product( ["two_kernel", "fused_atomic"], [8, 16], ["bytemap", "bitmap_u8", "dense_int16", "dense_int32", "dense_int64"], ) - for overwrite in ([False, True] if case[2].startswith("dense") else [True]) ] + [ - (execution, dtype, histogram_mode, topk, layout, False, True) + (execution, dtype, histogram_mode, topk, layout, (64, 32, 31)) for execution in ["eager", "compile"] for dtype, histogram_mode, topk, layout in [ (torch.float32, "two_kernel", 1, "bytemap"), @@ -852,27 +847,20 @@ def test_topk_sqrtsoftplus( ] ], ) -def test_qb_topk_histogram( - histogram_mode, - topk, - routing_output_mode, - overwrite, - execution, - dtype, - delta, -): - num_tokens, num_experts, num_bins = (64, 32, 31) if delta else (257, 896, 1000) - logits = torch.randn( - num_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True - ) - if delta: +def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, execution, dtype, shape): + num_tokens, num_experts, num_bins = shape + small = num_tokens == 64 + logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype, requires_grad=True) + if small: logits = ( _compile_logits(dtype).reshape(num_tokens, num_experts).detach().requires_grad_(True) ) reference_logits = logits.detach().float().requires_grad_(True) expert_bias = torch.linspace(-0.2, 0.2, num_experts, device="cuda", dtype=torch.float32) bin_bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) - reference_histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) + reference_histogram = torch.full( + (num_experts, num_bins), 7 if small else 0, device="cuda", dtype=torch.int32 + ) reference = qb_topk_score_function_pytorch( reference_logits, topk, @@ -883,7 +871,7 @@ def test_qb_topk_histogram( ) fused_logits = logits.detach().clone().requires_grad_(True) - fused_histogram = torch.full_like(reference_histogram, 7 if delta else 0) + fused_histogram = torch.full_like(reference_histogram, 7 if small else 0) histogram_version = fused_histogram._version dense_dtype = { "dense_int16": torch.int16, @@ -915,31 +903,21 @@ def test_qb_topk_histogram( qb_histogram=fused_histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, - overwrite_topk_indices=overwrite, - return_qb_histogram_delta=delta, - assume_valid_qb_bin_bounds=delta, + assume_valid_qb_bin_bounds=execution == "compile", ) fused_probs, fused_routing_output = outputs[:2] - if delta: + if small: _check_compile_reference(fused_probs, reference["probs"], dtype) else: torch.testing.assert_close(fused_probs, reference["probs"]) assert not fused_routing_output.requires_grad - if delta: - assert not outputs[2].requires_grad - assert outputs[2].data_ptr() != fused_histogram.data_ptr() - assert fused_histogram._version == histogram_version - torch.testing.assert_close(fused_histogram, torch.full_like(fused_histogram, 7)) - torch.testing.assert_close(outputs[2], reference["histogram"]) - fused_histogram.add_(1) - bin_bounds.add_(0.1) - else: - torch.testing.assert_close(fused_histogram, reference["histogram"]) + assert len(outputs) == 2 + assert fused_histogram._version > histogram_version + torch.testing.assert_close(fused_histogram, reference["histogram"]) + fused_histogram.add_(1) + bin_bounds.add_(0.1) if dense_dtype is not None: - assert (fused_routing_output.data_ptr() == topk_indices.data_ptr()) == overwrite - if not overwrite: - torch.testing.assert_close(topk_indices, torch.full_like(topk_indices, -1)) - topk_indices.fill_(-2) + assert fused_routing_output.data_ptr() == topk_indices.data_ptr() fused_routing_map = topk_indices_to_routing_map(fused_routing_output, num_experts) torch.testing.assert_close(fused_routing_map, reference["routing_map"]) elif routing_output_mode == "bitmap_u8": @@ -953,7 +931,7 @@ def test_qb_topk_histogram( grad = torch.randn((*fused_probs.shape, 2), device="cuda", dtype=dtype)[..., 0] reference["probs"].backward(grad.float()) fused_probs.backward(grad) - if delta: + if small: _check_compile_reference(fused_logits.grad, reference_logits.grad, dtype) else: torch.testing.assert_close(fused_logits.grad, reference_logits.grad, atol=1e-5, rtol=1e-5) @@ -969,17 +947,17 @@ def test_qb_topk_histogram( ], ) def test_qb_histogram_accumulates_microbatches(execution, dtype, histogram_mode): - functional = execution == "compile" - num_experts, num_bins = (130, 64) if functional else (64, 1000) + compiled = execution == "compile" + num_experts, num_bins = (130, 64) if compiled else (64, 1000) expert_bias = torch.linspace(-0.1, 0.1, num_experts, device="cuda") bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) histogram = torch.full( - (num_experts, num_bins), 7 if functional else 0, device="cuda", dtype=torch.int32 + (num_experts, num_bins), 7 if compiled else 0, device="cuda", dtype=torch.int32 ) reference_histogram = histogram.clone() - def route(logits, histogram, bounds): - outputs = fused_topk_with_score_function( + def route(logits, histogram, bounds, indices): + return fused_topk_with_score_function( logits, 8, False, @@ -988,41 +966,85 @@ def route(logits, histogram, bounds): None, "sigmoid", expert_bias, - routing_map_format="bitmap_u8" if functional else "bytemap", + routing_map_format="bitmap_u8" if compiled and indices is None else "bytemap", + topk_indices=indices, qb_histogram=histogram, qb_bin_bounds=bounds, qb_histogram_mode=histogram_mode, - return_qb_histogram_delta=functional, - assume_valid_qb_bin_bounds=functional, + assume_valid_qb_bin_bounds=compiled, ) - return outputs[0], outputs[1], histogram + outputs[2] if functional else histogram - if functional: + if compiled: torch._dynamo.reset() route = torch.compile(route, fullgraph=True, dynamic=True) - for num_tokens in (17, 33, 64) if functional else (127, 193): + pending = [] + for num_tokens in (17, 33, 64) if compiled else (127, 193): logits = torch.randn( num_tokens, num_experts, device="cuda", dtype=dtype, requires_grad=True ) reference = qb_topk_score_function_pytorch( logits, 8, expert_bias, bounds, num_bins, reference_histogram ) - input_histogram = histogram - probs, routing, histogram = route(logits, histogram, bounds) + indices = ( + torch.empty(num_tokens, 8, device="cuda", dtype=torch.int32) + if histogram_mode == "fused_atomic" + else None + ) + probs, routing = route(logits, histogram, bounds, indices) torch.testing.assert_close(histogram, reference["histogram"]) _check_compile_reference(probs, reference["probs"], dtype) - expected_map = ( - _bytemap_to_bitmap_u8(reference["routing_map"]) - if functional - else reference["routing_map"] - ) - torch.testing.assert_close(routing, expected_map) - if functional: - assert histogram.data_ptr() != input_histogram.data_ptr() - actual_grad = torch.autograd.grad(probs.float().square().sum(), logits)[0] - ref_grad = torch.autograd.grad(reference["probs"].float().square().sum(), logits)[0] - _check_compile_reference(actual_grad, ref_grad, dtype) + if indices is not None: + assert routing.data_ptr() == indices.data_ptr() + torch.testing.assert_close( + topk_indices_to_routing_map(routing, num_experts), reference["routing_map"] + ) + else: + expected_map = ( + _bytemap_to_bitmap_u8(reference["routing_map"]) + if compiled + else reference["routing_map"] + ) + torch.testing.assert_close(routing, expected_map) + pending.append((logits, probs, reference["probs"])) + if compiled: bounds.add_(0.03) + for logits, probs, reference_probs in reversed(pending): + actual_grad = torch.autograd.grad(probs.float().square().sum(), logits)[0] + ref_grad = torch.autograd.grad(reference_probs.float().square().sum(), logits)[0] + _check_compile_reference(actual_grad, ref_grad, dtype) + + +@pytest.mark.parametrize("execution", ["eager", "compile"]) +@pytest.mark.parametrize("histogram_mode", [None, "two_kernel", "fused_atomic"]) +def test_router_saved_indices(execution, histogram_mode): + logits = torch.randn(16, 32, device="cuda", requires_grad=True) + indices = torch.empty(16, 4, device="cuda", dtype=torch.int32) + bias = torch.zeros(32, device="cuda") + histogram = torch.zeros(32, 31, device="cuda", dtype=torch.int32) + bounds = torch.tensor([-1.0, 1.0], device="cuda") + + def route(x): + return fused_topk_with_score_function( + x, + 4, + False, + None, + None, + None, + "sigmoid", + bias, + topk_indices=indices, + qb_histogram=histogram if histogram_mode else None, + qb_bin_bounds=bounds if histogram_mode else None, + qb_histogram_mode=histogram_mode, + assume_valid_qb_bin_bounds=True, + ) + + route = _router_function(route, execution) + probs, _ = route(logits) + route(-logits) + with pytest.raises(RuntimeError, match="modified by an inplace operation"): + probs.sum().backward() def test_qb_topk_argument_validation(): @@ -1057,19 +1079,6 @@ def test_qb_topk_argument_validation(): qb_histogram_mode="two_kernel", ) - with pytest.raises(ValueError, match="return_qb_histogram_delta requires"): - fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - None, - return_qb_histogram_delta=True, - ) - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least two CUDA devices") @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @@ -1106,8 +1115,7 @@ def test_qb_topk_uses_logits_device(histogram_mode): @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @pytest.mark.parametrize("invalid_bounds", ["equal", "reversed", "nonfinite"]) -@pytest.mark.parametrize("delta", [False, True]) -def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds, delta): +def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds): logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) @@ -1130,7 +1138,6 @@ def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds, delt qb_histogram=histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, - return_qb_histogram_delta=delta, ) @@ -1536,7 +1543,6 @@ def test_topk_preserves_leading_dims(topk_index_dtype, execution): score_function="softmax", expert_bias=None, topk_indices=topk_indices, - overwrite_topk_indices=execution == "eager", ) assert probs.shape == logits.shape diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 7d0e216f121..e208dfe6520 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -160,20 +160,35 @@ def forward( topk_indices: Optional[torch.Tensor], ): # pylint: disable=missing-function-docstring - probs, routing_output, intermediate_output = tex.fused_topk_with_score_function_fwd( - logits, - topk, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - score_function, - expert_bias, - routing_map_format, - topk_indices, - ) - if topk_indices is not None: - routing_output = topk_indices + ctx.compiled = torch.compiler.is_compiling() + if ctx.compiled: + probs, routing_output, intermediate_output = _topk_forward( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, + ) + else: + probs, routing_output, intermediate_output = tex.fused_topk_with_score_function_fwd( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, + ) + if topk_indices is not None: + routing_output = topk_indices if topk_indices is not None: ctx.mark_dirty(topk_indices) ctx.mark_non_differentiable(routing_output) @@ -190,6 +205,19 @@ def forward( def backward(ctx, grad_probs, _): # pylint: disable=missing-function-docstring routing_map, intermediate_output = ctx.saved_tensors + if ctx.compiled: + grad_logits = _topk_backward( + routing_map, + intermediate_output, + grad_probs, + ctx.topk, + ctx.use_pre_softmax, + ctx.scaling_factor, + ctx.score_function, + ctx.routing_map_format, + ctx.use_dense_indices, + ) + return grad_logits, None, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -226,29 +254,45 @@ def forward( assume_valid_bin_bounds: bool, ): # pylint: disable=missing-function-docstring - bin_bounds_validated = assume_valid_bin_bounds or _validate_qb_bin_bounds(bin_bounds) - ( - probs, - routing_output, - intermediate_output, - _cutoff, - histogram_output, - ) = tex.fused_topk_with_score_function_qb_fwd( - logits, - topk, - scaling_factor, - expert_bias, - routing_map_format, - topk_indices, - histogram, - bin_bounds, - histogram_mode, - bin_bounds_validated, - ) + ctx.compiled = torch.compiler.is_compiling() + if ctx.compiled: + probs, routing_output, intermediate_output = _qb_forward( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + ) + histogram_output = histogram + else: + bin_bounds_validated = assume_valid_bin_bounds or _validate_qb_bin_bounds(bin_bounds) + ( + probs, + routing_output, + intermediate_output, + _cutoff, + histogram_output, + ) = tex.fused_topk_with_score_function_qb_fwd( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + bin_bounds_validated, + ) if topk_indices is not None: routing_output = topk_indices - ctx.mark_dirty(topk_indices) - ctx.mark_dirty(histogram) + ctx.mark_dirty(topk_indices, histogram) + else: + ctx.mark_dirty(histogram) ctx.mark_non_differentiable(routing_output, histogram_output) ctx.save_for_backward(routing_output, intermediate_output) ctx.topk = topk @@ -261,6 +305,19 @@ def forward( def backward(ctx, grad_probs, _, _histogram_grad): # pylint: disable=missing-function-docstring routing_output, intermediate_output = ctx.saved_tensors + if ctx.compiled: + grad_logits = _topk_backward( + routing_output, + intermediate_output, + grad_probs, + ctx.topk, + False, + ctx.scaling_factor, + "sigmoid", + ctx.routing_map_format, + ctx.use_dense_indices, + ) + return grad_logits, None, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -280,14 +337,10 @@ def backward(ctx, grad_probs, _, _histogram_grad): def _needs_eager_router(call): - if call.get("topk_indices") is not None and call.get("overwrite_topk_indices", True): - return "router output buffers" if any( call.get(name) is not None for name in ("qb_histogram", "qb_bin_bounds", "qb_histogram_mode") ): - if not call.get("return_qb_histogram_delta", False): - return "in-place QB histogram accumulation" if not call.get("assume_valid_qb_bin_bounds", False): return "QB bin-bound value validation" return None @@ -308,18 +361,16 @@ def fused_topk_with_score_function( qb_histogram: Optional[torch.Tensor] = None, qb_bin_bounds: Optional[torch.Tensor] = None, qb_histogram_mode: Optional[str] = None, - overwrite_topk_indices: bool = True, - return_qb_histogram_delta: bool = False, assume_valid_qb_bin_bounds: bool = False, ): """ Fused topk with score function router. - ``torch.compile`` supports routing maps and dense indices with - ``overwrite_topk_indices=False``. Quantile Balancing also requires - ``return_qb_histogram_delta=True`` and ``assume_valid_qb_bin_bounds=True``. - Buffer mutation and QB value validation fall back to eager execution and do not - support ``fullgraph=True``. + ``torch.compile(fullgraph=True)`` supports routing maps, dense output buffers, and + in-place QB histogram accumulation. QB requires ``assume_valid_qb_bin_bounds=True``; + otherwise, bounds value validation falls back to eager execution. For automatic CUDA + graph replay, use module buffers or buffers marked with ``torch._dynamo.mark_static_address`` + for mutated inputs. Ordinary mutated inputs may cause CUDA graph capture to be skipped. Parameters ---------- @@ -345,11 +396,10 @@ def fused_topk_with_score_function( topk_indices : torch.Tensor, optional Optional output buffer with shape [*logits.shape[:-1], topk]. When provided, its dtype controls the dense index output dtype and the routing map is not materialized. - With ``overwrite_topk_indices=False``, only its metadata is used. + The buffer is overwritten and returned in both eager and compiled execution. + Its contents must remain unchanged until the corresponding backward completes. qb_histogram : torch.Tensor, optional Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. - With ``return_qb_histogram_delta=True``, only its metadata is used; the supplied - histogram is neither read nor modified. qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be finite with ``lower < upper``. With default value validation, eager calls revalidate @@ -362,13 +412,6 @@ def fused_topk_with_score_function( Invalid replay-time bounds may silently produce an incorrect histogram. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. - overwrite_topk_indices : bool, default = True - Write indices into ``topk_indices`` and return that buffer. If False, return a new - tensor without modifying the supplied buffer. Ignored when ``topk_indices`` is None. - return_qb_histogram_delta : bool, default = False - Return the current call's histogram as a third output instead of accumulating into - ``qb_histogram``. Requires all three QB arguments. The caller owns accumulation, - for example ``histogram = histogram + histogram_delta`` across microbatches. assume_valid_qb_bin_bounds : bool, default = False Skip QB bounds value validation and its CPU synchronization. The caller must guarantee finite bounds with ``lower < upper`` on every call and CUDA graph replay, @@ -386,20 +429,14 @@ def fused_topk_with_score_function( - BITMAP_U8: uint8[*logits.shape[:-1], ceil(num_experts/8)] LSB-first bit-packed. - Dense indices: [*logits.shape[:-1], topk], with the dtype of ``topk_indices``. - histogram_delta : torch.Tensor, optional - Returned only when ``return_qb_histogram_delta=True``. A new non-differentiable - int32 ``[num_experts, num_bins]`` tensor containing this call's counts, excluding - all previous counts in ``qb_histogram``. """ if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type.") routing_map_format = _validate_routing_map_format(routing_map_format) - if ( - not torch.compiler.is_compiling() - and topk_indices is not None - and not overwrite_topk_indices - ): - topk_indices = torch.empty_like(topk_indices) + compiled_inference = torch.compiler.is_compiling() and ( + not torch.is_grad_enabled() + or not any(t is not None and t.requires_grad for t in (logits, expert_bias, qb_bin_bounds)) + ) qb_arguments = (qb_histogram, qb_bin_bounds, qb_histogram_mode) if any(value is not None for value in qb_arguments): if any(value is None for value in qb_arguments): @@ -418,8 +455,8 @@ def fused_topk_with_score_function( "qb_histogram_mode must be 'two_kernel' or 'fused_atomic', " f"got {qb_histogram_mode!r}" ) - if torch.compiler.is_compiling(): - probs, routing_output, _, histogram_delta = _qb_forward( + if compiled_inference: + probs, routing_output, _ = _qb_forward( logits, topk, scaling_factor, @@ -430,10 +467,8 @@ def fused_topk_with_score_function( qb_bin_bounds, mode, ) - return probs, routing_output, histogram_delta - if return_qb_histogram_delta: - qb_histogram = torch.zeros_like(qb_histogram) - probs, routing_output, histogram_output = FusedTopkScoreFunctionQB.apply( + return probs, routing_output + probs, routing_output, _ = FusedTopkScoreFunctionQB.apply( logits, topk, scaling_factor, @@ -445,12 +480,8 @@ def fused_topk_with_score_function( mode, assume_valid_qb_bin_bounds, ) - if return_qb_histogram_delta: - return probs, routing_output, histogram_output return probs, routing_output - if return_qb_histogram_delta: - raise ValueError("return_qb_histogram_delta requires all three QB arguments") - if torch.compiler.is_compiling(): + if compiled_inference: probs, routing_output, _ = _topk_forward( logits, topk, @@ -670,8 +701,83 @@ def _routing_map_fake(logits, routing_map_format): return logits.new_empty(logits.shape, dtype=torch.bool) -@torch.library.custom_op("te_moe::router_topk_fwd", mutates_args=()) def _topk_forward( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, +): + if topk_indices is None: + return _topk_map_forward( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + ) + probs, intermediate_output = _topk_indices_forward( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, + ) + return probs, topk_indices, intermediate_output + + +def _qb_forward( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, +): + if topk_indices is None: + return _qb_map_forward( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + histogram, + bin_bounds, + histogram_mode, + ) + probs, intermediate_output = _qb_indices_forward( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + ) + return probs, topk_indices, intermediate_output + + +@torch.library.custom_op("te_moe::router_topk_map_fwd", mutates_args=()) +def _topk_map_forward( logits: torch.Tensor, topk: int, use_pre_softmax: bool, @@ -681,10 +787,8 @@ def _topk_forward( score_function: str, expert_bias: Optional[torch.Tensor], routing_map_format: int, - topk_indices: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: with torch.cuda.device(logits.device): - indices = torch.empty_like(topk_indices) if topk_indices is not None else None return tex.fused_topk_with_score_function_fwd( logits, topk, @@ -695,12 +799,12 @@ def _topk_forward( score_function, expert_bias, routing_map_format, - indices, + None, ) -@_topk_forward.register_fake -def _topk_forward_fake( +@_topk_map_forward.register_fake +def _topk_map_forward_fake( logits, topk, use_pre_softmax, @@ -710,7 +814,6 @@ def _topk_forward_fake( score_function, expert_bias, routing_map_format, - topk_indices=None, ): del ( topk, @@ -723,15 +826,58 @@ def _topk_forward_fake( ) return ( logits.new_empty(logits.shape), - ( - torch.empty_like(topk_indices) - if topk_indices is not None - else _routing_map_fake(logits, routing_map_format) - ), + _routing_map_fake(logits, routing_map_format), logits.new_empty(logits.shape, dtype=torch.float32), ) +@torch.library.custom_op("te_moe::router_topk_indices_fwd", mutates_args=("topk_indices",)) +def _topk_indices_forward( + logits: torch.Tensor, + topk: int, + use_pre_softmax: bool, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, + expert_bias: Optional[torch.Tensor], + routing_map_format: int, + topk_indices: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + with torch.cuda.device(logits.device): + probs, _, intermediate_output = tex.fused_topk_with_score_function_fwd( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, + ) + return probs, intermediate_output + + +@_topk_indices_forward.register_fake +def _topk_indices_forward_fake( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, +): + del topk, use_pre_softmax, num_groups, group_topk, scaling_factor + del score_function, expert_bias, routing_map_format, topk_indices + return torch.empty_like(logits), torch.empty_like(logits, dtype=torch.float32) + + @torch.library.custom_op("te_moe::router_topk_bwd", mutates_args=()) def _topk_backward( routing_map: torch.Tensor, @@ -779,126 +925,99 @@ def _topk_backward_fake( return grad_probs.new_empty(grad_probs.shape) -def _topk_setup_context(ctx, inputs, output): - ( - _, - ctx.topk, - ctx.use_pre_softmax, - _, - _, - ctx.scaling_factor, - ctx.score_function, - _, - ctx.routing_map_format, - topk_indices, - ) = inputs - ctx.use_dense_indices = topk_indices is not None - _, routing_map, intermediate_output = output - ctx.mark_non_differentiable(routing_map, intermediate_output) - ctx.save_for_backward(routing_map, intermediate_output) - - -def _topk_autograd(ctx, grad_probs, _grad_map, _grad_intermediate): - routing_map, intermediate_output = ctx.saved_tensors - grad_logits = _topk_backward( - routing_map, - intermediate_output, - grad_probs, - ctx.topk, - ctx.use_pre_softmax, - ctx.scaling_factor, - ctx.score_function, - ctx.routing_map_format, - ctx.use_dense_indices, - ) - return grad_logits, None, None, None, None, None, None, None, None, None - - -_topk_forward.register_autograd(_topk_autograd, setup_context=_topk_setup_context) - - -@torch.library.custom_op("te_moe::router_qb_fwd", mutates_args=()) -def _qb_forward( +@torch.library.custom_op("te_moe::router_qb_map_fwd", mutates_args=("histogram",)) +def _qb_map_forward( logits: torch.Tensor, topk: int, scaling_factor: Optional[float], expert_bias: torch.Tensor, routing_map_format: int, - topk_indices: Optional[torch.Tensor], histogram: torch.Tensor, bin_bounds: torch.Tensor, histogram_mode: int, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: with torch.cuda.device(logits.device): - indices = torch.empty_like(topk_indices) if topk_indices is not None else None - histogram_delta = torch.zeros_like(histogram) - probs, routing_output, intermediate_output, _, histogram_delta = ( + probs, routing_output, intermediate_output, _, _ = ( tex.fused_topk_with_score_function_qb_fwd( logits, topk, scaling_factor, expert_bias, routing_map_format, - indices, - histogram_delta, + None, + histogram, bin_bounds, histogram_mode, True, ) ) - return probs, routing_output, intermediate_output, histogram_delta + return probs, routing_output, intermediate_output -@_qb_forward.register_fake -def _qb_forward_fake( +@_qb_map_forward.register_fake +def _qb_map_forward_fake( logits, topk, scaling_factor, expert_bias, routing_map_format, - topk_indices, histogram, bin_bounds, histogram_mode, ): - del topk, scaling_factor, expert_bias, bin_bounds, histogram_mode + del topk, scaling_factor, expert_bias, histogram, bin_bounds, histogram_mode return ( logits.new_empty(logits.shape), - ( - torch.empty_like(topk_indices) - if topk_indices is not None - else _routing_map_fake(logits, routing_map_format) - ), + _routing_map_fake(logits, routing_map_format), logits.new_empty(logits.shape, dtype=torch.float32), - torch.empty_like(histogram), ) -def _qb_setup_context(ctx, inputs, output): - _, ctx.topk, ctx.scaling_factor, _, ctx.routing_map_format, indices, _, _, _ = inputs - ctx.use_dense_indices = indices is not None - _, routing_output, intermediate_output, histogram_delta = output - ctx.mark_non_differentiable(routing_output, intermediate_output, histogram_delta) - ctx.save_for_backward(routing_output, intermediate_output) - - -def _qb_autograd(ctx, grad_probs, _grad_routing, _grad_intermediate, _grad_histogram): - routing_output, intermediate_output = ctx.saved_tensors - grad_logits = _topk_backward( - routing_output, - intermediate_output, - grad_probs, - ctx.topk, - False, - ctx.scaling_factor, - "sigmoid", - ctx.routing_map_format, - ctx.use_dense_indices, - ) - return grad_logits, None, None, None, None, None, None, None, None +@torch.library.custom_op( + "te_moe::router_qb_indices_fwd", mutates_args=("topk_indices", "histogram") +) +def _qb_indices_forward( + logits: torch.Tensor, + topk: int, + scaling_factor: Optional[float], + expert_bias: torch.Tensor, + routing_map_format: int, + topk_indices: torch.Tensor, + histogram: torch.Tensor, + bin_bounds: torch.Tensor, + histogram_mode: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + with torch.cuda.device(logits.device): + probs, _, intermediate_output, _, _ = tex.fused_topk_with_score_function_qb_fwd( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + True, + ) + return probs, intermediate_output -_qb_forward.register_autograd(_qb_autograd, setup_context=_qb_setup_context) +@_qb_indices_forward.register_fake +def _qb_indices_forward_fake( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, +): + del topk, scaling_factor, expert_bias, routing_map_format, topk_indices + del histogram, bin_bounds, histogram_mode + return torch.empty_like(logits), torch.empty_like(logits, dtype=torch.float32) @torch.library.custom_op("te_moe::router_score_fwd", mutates_args=()) From 9f24dfc48137f172f4504e9ac2cd96b97954d71a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 15:45:40 +0200 Subject: [PATCH 06/14] Skip QB bounds value validation under torch.compile Remove the opt-in flag and eager fallback for compiled QB routing. Keep metadata checks and eager version-cached value validation, and document the compiled caller contract for bounds updates and CUDA graph replay. Validate fullgraph training and inference without host value checks, retain eager invalid-bounds errors, and update the existing compile tests. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 74 ++++++++-------------------- transformer_engine/pytorch/router.py | 44 +++++------------ 2 files changed, 34 insertions(+), 84 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index f6776af3ebd..66fb8fe667c 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -173,7 +173,6 @@ def router(logits, bounds, total, indices, histogram): qb_histogram=histogram if histogram_mode else None, qb_bin_bounds=bounds if histogram_mode else None, qb_histogram_mode=histogram_mode, - assume_valid_qb_bin_bounds=histogram_mode is not None, ) _, scores = fused_compute_score_for_moe_aux_loss(logits, 4, "sigmoid") loss = fused_moe_aux_loss(scores.reshape(-1, 32), counts, total, 32, 4, 0.01) @@ -328,16 +327,18 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogra ) -@pytest.mark.parametrize("index_dtype", [None, torch.int32]) -def test_router_compile_eager_fallback(index_dtype): +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +@pytest.mark.parametrize("requires_grad", [False, True]) +def test_qb_compile_skips_value_validation(monkeypatch, histogram_mode, requires_grad): + from transformer_engine.pytorch import router + torch._dynamo.reset() - logits = torch.randn(32, 32, device="cuda", requires_grad=True) - bias = torch.zeros(32, device="cuda") - indices = torch.empty(32, 4, device="cuda", dtype=torch.int32) - histogram = torch.zeros(32, 16, device="cuda", dtype=torch.int32) - bounds = torch.tensor([-2, 2], device="cuda", dtype=torch.float32) + logits = torch.randn(8, 16, device="cuda", requires_grad=requires_grad) + bias = torch.zeros(16, device="cuda") + bounds = torch.tensor([-1.0, 1.0], device="cuda") + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - def router(x): + def route(x): return fused_topk_with_score_function( x, 4, @@ -347,60 +348,31 @@ def router(x): None, "sigmoid", bias, - topk_indices=indices if index_dtype is not None else None, qb_histogram=histogram, qb_bin_bounds=bounds, - qb_histogram_mode="two_kernel", + qb_histogram_mode=histogram_mode, ) - with pytest.warns(UserWarning, match="Falling back to eager execution"): - actual = torch.compile(router)(logits) - expected = router(logits) - torch.testing.assert_close(actual, expected) - torch._dynamo.reset() - with pytest.raises(torch._dynamo.exc.Unsupported): - torch.compile(router, fullgraph=True)(logits) - - -def test_qb_trusted_bounds_skip_value_validation(monkeypatch): - from transformer_engine.pytorch import router - - logits = torch.randn(8, 16, device="cuda", requires_grad=True) - bias = torch.zeros(16, device="cuda") - bounds = torch.tensor([-1.0, 1.0], device="cuda") - options = dict( - qb_histogram=torch.zeros(16, 32, device="cuda", dtype=torch.int32), - qb_bin_bounds=bounds, - qb_histogram_mode="two_kernel", - ) - def unexpected_validation(_): - raise AssertionError("trusted QB bounds must not read values on the CPU") + raise AssertionError("compiled QB must not read bounds values on the CPU") + compiled = torch.compile(route, fullgraph=True) with monkeypatch.context() as patch: patch.setattr(router, "_validate_qb_bin_bounds", unexpected_validation) - fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - bias, - **options, - assume_valid_qb_bin_bounds=True, - ) + for _ in range(2): + probs, _ = compiled(logits) + if requires_grad: + probs.square().sum().backward() + bounds.add_(0.1) + torch.testing.assert_close(histogram.sum(dim=-1), torch.full((16,), 16, device="cuda")) bounds.zero_() with pytest.raises(ValueError, match="finite with lower < upper"): - fused_topk_with_score_function( - logits, 4, False, None, None, None, "sigmoid", bias, **options - ) + route(logits) @pytest.mark.parametrize("compiled", [False, True]) @pytest.mark.parametrize("invalid", ["shape", "dtype", "device"]) -def test_qb_trusted_bounds_check_metadata(compiled, invalid): +def test_qb_bounds_check_metadata(compiled, invalid): torch._dynamo.reset() logits = torch.randn(8, 16, device="cuda") bias = torch.zeros(16, device="cuda") @@ -426,7 +398,6 @@ def router(logits, bounds): qb_histogram=histogram, qb_bin_bounds=bounds, qb_histogram_mode="two_kernel", - assume_valid_qb_bin_bounds=True, ) if compiled: @@ -903,7 +874,6 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, execution, qb_histogram=fused_histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, - assume_valid_qb_bin_bounds=execution == "compile", ) fused_probs, fused_routing_output = outputs[:2] if small: @@ -971,7 +941,6 @@ def route(logits, histogram, bounds, indices): qb_histogram=histogram, qb_bin_bounds=bounds, qb_histogram_mode=histogram_mode, - assume_valid_qb_bin_bounds=compiled, ) if compiled: @@ -1037,7 +1006,6 @@ def route(x): qb_histogram=histogram if histogram_mode else None, qb_bin_bounds=bounds if histogram_mode else None, qb_histogram_mode=histogram_mode, - assume_valid_qb_bin_bounds=True, ) route = _router_function(route, execution) diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index e208dfe6520..7e82f143126 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -20,8 +20,6 @@ import torch import transformer_engine_torch as tex -from .jit import no_torch_dynamo - # Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. # Members: # RoutingMapFormat.BYTEMAP — bool[num_tokens, num_experts] @@ -251,7 +249,6 @@ def forward( histogram: torch.Tensor, bin_bounds: torch.Tensor, histogram_mode: int, - assume_valid_bin_bounds: bool, ): # pylint: disable=missing-function-docstring ctx.compiled = torch.compiler.is_compiling() @@ -269,7 +266,7 @@ def forward( ) histogram_output = histogram else: - bin_bounds_validated = assume_valid_bin_bounds or _validate_qb_bin_bounds(bin_bounds) + bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -317,7 +314,7 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.routing_map_format, ctx.use_dense_indices, ) - return grad_logits, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -333,20 +330,9 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.use_dense_indices, ctx.routing_map_format, ) - return grad_logits, None, None, None, None, None, None, None, None, None - - -def _needs_eager_router(call): - if any( - call.get(name) is not None - for name in ("qb_histogram", "qb_bin_bounds", "qb_histogram_mode") - ): - if not call.get("assume_valid_qb_bin_bounds", False): - return "QB bin-bound value validation" - return None + return grad_logits, None, None, None, None, None, None, None, None -@no_torch_dynamo(when=_needs_eager_router) def fused_topk_with_score_function( logits: torch.Tensor, topk: int, @@ -361,15 +347,14 @@ def fused_topk_with_score_function( qb_histogram: Optional[torch.Tensor] = None, qb_bin_bounds: Optional[torch.Tensor] = None, qb_histogram_mode: Optional[str] = None, - assume_valid_qb_bin_bounds: bool = False, ): """ Fused topk with score function router. ``torch.compile(fullgraph=True)`` supports routing maps, dense output buffers, and - in-place QB histogram accumulation. QB requires ``assume_valid_qb_bin_bounds=True``; - otherwise, bounds value validation falls back to eager execution. For automatic CUDA - graph replay, use module buffers or buffers marked with ``torch._dynamo.mark_static_address`` + in-place QB histogram accumulation. Compiled execution skips QB bounds value validation + and its CPU synchronization; tensor metadata is still checked. For automatic CUDA graph + replay, use module buffers or buffers marked with ``torch._dynamo.mark_static_address`` for mutated inputs. Ordinary mutated inputs may cause CUDA graph capture to be skipped. Parameters @@ -402,21 +387,19 @@ def fused_topk_with_score_function( Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be - finite with ``lower < upper``. With default value validation, eager calls revalidate - PyTorch-tracked in-place updates. Before eager CUDA graph capture, validate the - current version eagerly or use + finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. + Under ``torch.compile``, the caller must guarantee valid values on every call, + including after device-side updates; compiled calls do not mark bounds as validated + for later eager calls. Before eager CUDA graph capture, validate the current version + eagerly or use :func:`mark_qb_bin_bounds_validated` after a trusted device-side update. CUDA graph replay does not execute Python validation or advance PyTorch tensor-version counters. Any bounds update visible to replay, whether issued outside the graph or captured inside it, is therefore trusted and must produce finite values with ``lower < upper`` on every replay. - Invalid replay-time bounds may silently produce an incorrect histogram. + Invalid bounds under ``torch.compile`` or CUDA graph replay may silently produce an + incorrect histogram. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. - assume_valid_qb_bin_bounds : bool, default = False - Skip QB bounds value validation and its CPU synchronization. The caller must - guarantee finite bounds with ``lower < upper`` on every call and CUDA graph replay, - including after device-side updates. Tensor metadata is still checked. Ignored - without QB. Setting this does not mark the tensor as validated for later checked calls. Returns ------- @@ -478,7 +461,6 @@ def fused_topk_with_score_function( qb_histogram, qb_bin_bounds, mode, - assume_valid_qb_bin_bounds, ) return probs, routing_output if compiled_inference: From a3a6a10ca3e57001a6ef1cb709d52086aa80d256 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 16:07:58 +0200 Subject: [PATCH 07/14] Simplify compiled router inference context handling Use the same autograd.Function entry point for training and inference. Pass the compiled-inference decision from the wrapper and return before context setup, avoiding Dynamo's unsupported inference mark_dirty path. Keep training context and eager native dispatch unchanged. Exercise no_grad, inference_mode, and inputs without requires_grad in the existing CUDA graph replay tests. Validation: 3391 router tests passed, 446 skipped; pre-commit, license, and router pylint passed. Full Python lint retains the existing C0411 in fused_mla_q_uproj.py. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 22 ++++++++---- transformer_engine/pytorch/router.py | 53 ++++++++++------------------ 2 files changed, 34 insertions(+), 41 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 66fb8fe667c..eff3704f8a0 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -205,12 +205,22 @@ def router(logits, bounds, total, indices, histogram): actual_grad = torch.autograd.grad(actual[0], logits) expected_grad = torch.autograd.grad(expected[0], eager_logits) torch.testing.assert_close(actual_grad, expected_grad) - with torch.no_grad(): - torch.compiler.cudagraph_mark_step_begin() - actual = compiled(logits, bounds, total, indices, histogram) - expected = router(logits, bounds, total, reference_indices, reference_histogram) - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(histogram, reference_histogram) + for grad_context, inference_logits in ( + (torch.no_grad, logits), + (torch.inference_mode, logits), + (torch.enable_grad, logits.detach()), + ): + with grad_context(): + torch.compiler.cudagraph_mark_step_begin() + actual = compiled(inference_logits, bounds, total, indices, histogram) + expected = router( + inference_logits, bounds, total, reference_indices, reference_histogram + ) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(histogram, reference_histogram) + assert not actual[0].requires_grad + if indices is not None: + assert actual[1].data_ptr() == indices.data_ptr() assert not counters["inductor"]["cudagraph_skips"] diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 7e82f143126..b2fcaaf9b37 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -156,10 +156,11 @@ def forward( expert_bias: Optional[torch.Tensor], routing_map_format: int, topk_indices: Optional[torch.Tensor], + skip_context: bool, ): # pylint: disable=missing-function-docstring - ctx.compiled = torch.compiler.is_compiling() - if ctx.compiled: + compiled = torch.compiler.is_compiling() + if compiled: probs, routing_output, intermediate_output = _topk_forward( logits, topk, @@ -172,6 +173,8 @@ def forward( routing_map_format, topk_indices, ) + if skip_context: + return probs, routing_output else: probs, routing_output, intermediate_output = tex.fused_topk_with_score_function_fwd( logits, @@ -187,6 +190,7 @@ def forward( ) if topk_indices is not None: routing_output = topk_indices + ctx.compiled = compiled if topk_indices is not None: ctx.mark_dirty(topk_indices) ctx.mark_non_differentiable(routing_output) @@ -215,7 +219,7 @@ def backward(ctx, grad_probs, _): ctx.routing_map_format, ctx.use_dense_indices, ) - return grad_logits, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -231,7 +235,7 @@ def backward(ctx, grad_probs, _): ctx.use_dense_indices, ctx.routing_map_format, ) - return grad_logits, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None, None class FusedTopkScoreFunctionQB(torch.autograd.Function): @@ -249,10 +253,11 @@ def forward( histogram: torch.Tensor, bin_bounds: torch.Tensor, histogram_mode: int, + skip_context: bool, ): # pylint: disable=missing-function-docstring - ctx.compiled = torch.compiler.is_compiling() - if ctx.compiled: + compiled = torch.compiler.is_compiling() + if compiled: probs, routing_output, intermediate_output = _qb_forward( logits, topk, @@ -264,6 +269,8 @@ def forward( bin_bounds, histogram_mode, ) + if skip_context: + return probs, routing_output, histogram histogram_output = histogram else: bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) @@ -285,6 +292,7 @@ def forward( histogram_mode, bin_bounds_validated, ) + ctx.compiled = compiled if topk_indices is not None: routing_output = topk_indices ctx.mark_dirty(topk_indices, histogram) @@ -314,7 +322,7 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.routing_map_format, ctx.use_dense_indices, ) - return grad_logits, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -330,7 +338,7 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.use_dense_indices, ctx.routing_map_format, ) - return grad_logits, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None def fused_topk_with_score_function( @@ -438,19 +446,6 @@ def fused_topk_with_score_function( "qb_histogram_mode must be 'two_kernel' or 'fused_atomic', " f"got {qb_histogram_mode!r}" ) - if compiled_inference: - probs, routing_output, _ = _qb_forward( - logits, - topk, - scaling_factor, - expert_bias, - routing_map_format, - topk_indices, - qb_histogram, - qb_bin_bounds, - mode, - ) - return probs, routing_output probs, routing_output, _ = FusedTopkScoreFunctionQB.apply( logits, topk, @@ -461,20 +456,7 @@ def fused_topk_with_score_function( qb_histogram, qb_bin_bounds, mode, - ) - return probs, routing_output - if compiled_inference: - probs, routing_output, _ = _topk_forward( - logits, - topk, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - score_function, - expert_bias, - routing_map_format, - topk_indices, + compiled_inference, ) return probs, routing_output return FusedTopkScoreFunction.apply( @@ -488,6 +470,7 @@ def fused_topk_with_score_function( expert_bias, routing_map_format, topk_indices, + compiled_inference, ) From ea3275cb79b22ffc5efeeda3599a067840e9da51 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 18:52:28 +0200 Subject: [PATCH 08/14] Remove unrelated router test formatting change Restore the existing reference-call layout to keep the compile PR focused. The complete test module AST is unchanged. Black with repository flags and git diff --check pass. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index eff3704f8a0..27357bf27bb 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -43,11 +43,7 @@ def test_sigmoid_bias_preserves_tiny_scores(topk, histogram_mode, layout, mixed_ bin_bounds = torch.tensor([-1.0, 1.0], device="cuda") histogram = torch.zeros(num_experts, num_bins, dtype=torch.int32, device="cuda") reference_probs, reference_map = topk_score_function_pytorch( - logits, - topk, - scaling_factor=3.16, - score_function="sigmoid", - expert_bias=expert_bias, + logits, topk, scaling_factor=3.16, score_function="sigmoid", expert_bias=expert_bias ) kwargs = {} if histogram_mode is not None: From 41219e9ea7cd88f3bfef3509357e260b7d7c1f77 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 19:26:01 +0200 Subject: [PATCH 09/14] Simplify fused router tests and helper names Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 52 ++++++------------------------ 1 file changed, 9 insertions(+), 43 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 27357bf27bb..92779b05ea9 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -121,13 +121,13 @@ def _router_function(fn, execution): return torch.compile(fn, fullgraph=True) -def _compile_logits(dtype): +def _make_logits(dtype): logits = torch.linspace(-2, 2, 32, device="cuda") logits = logits + torch.linspace(-0.1, 0.1, 64, device="cuda").unsqueeze(1) return logits.reshape(2, 32, 32).to(dtype).requires_grad_(True) -def _check_compile_reference(actual, expected, dtype): +def _check_reference(actual, expected, dtype): atol, rtol = { torch.float32: (1e-6, 2e-5), torch.float16: (2e-4, 2e-3), @@ -184,7 +184,7 @@ def router(logits, bounds, total, indices, histogram): torch.compiler.cudagraph_mark_step_begin() bounds.copy_(bounds.new_tensor(new_bounds)) total.fill_(64 + 17 * step) - logits = _compile_logits(torch.bfloat16) + logits = _make_logits(torch.bfloat16) if step == 1: logits = logits.flip(-1).detach().requires_grad_(True) eager_logits = logits.detach().clone().requires_grad_(True) @@ -261,7 +261,7 @@ def router_loss(logits): torch.testing.assert_close(actual, expected) (actual_grad,) = torch.autograd.grad(actual, logits) (eager_grad,) = torch.autograd.grad(expected, eager_logits) - _check_compile_reference(actual_grad, eager_grad, dtype) + _check_reference(actual_grad, eager_grad, dtype) @pytest.mark.parametrize( @@ -829,9 +829,7 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, execution, small = num_tokens == 64 logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype, requires_grad=True) if small: - logits = ( - _compile_logits(dtype).reshape(num_tokens, num_experts).detach().requires_grad_(True) - ) + logits = _make_logits(dtype).reshape(num_tokens, num_experts).detach().requires_grad_(True) reference_logits = logits.detach().float().requires_grad_(True) expert_bias = torch.linspace(-0.2, 0.2, num_experts, device="cuda", dtype=torch.float32) bin_bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) @@ -883,7 +881,7 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, execution, ) fused_probs, fused_routing_output = outputs[:2] if small: - _check_compile_reference(fused_probs, reference["probs"], dtype) + _check_reference(fused_probs, reference["probs"], dtype) else: torch.testing.assert_close(fused_probs, reference["probs"]) assert not fused_routing_output.requires_grad @@ -908,7 +906,7 @@ def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode, execution, reference["probs"].backward(grad.float()) fused_probs.backward(grad) if small: - _check_compile_reference(fused_logits.grad, reference_logits.grad, dtype) + _check_reference(fused_logits.grad, reference_logits.grad, dtype) else: torch.testing.assert_close(fused_logits.grad, reference_logits.grad, atol=1e-5, rtol=1e-5) @@ -967,7 +965,7 @@ def route(logits, histogram, bounds, indices): ) probs, routing = route(logits, histogram, bounds, indices) torch.testing.assert_close(histogram, reference["histogram"]) - _check_compile_reference(probs, reference["probs"], dtype) + _check_reference(probs, reference["probs"], dtype) if indices is not None: assert routing.data_ptr() == indices.data_ptr() torch.testing.assert_close( @@ -986,39 +984,7 @@ def route(logits, histogram, bounds, indices): for logits, probs, reference_probs in reversed(pending): actual_grad = torch.autograd.grad(probs.float().square().sum(), logits)[0] ref_grad = torch.autograd.grad(reference_probs.float().square().sum(), logits)[0] - _check_compile_reference(actual_grad, ref_grad, dtype) - - -@pytest.mark.parametrize("execution", ["eager", "compile"]) -@pytest.mark.parametrize("histogram_mode", [None, "two_kernel", "fused_atomic"]) -def test_router_saved_indices(execution, histogram_mode): - logits = torch.randn(16, 32, device="cuda", requires_grad=True) - indices = torch.empty(16, 4, device="cuda", dtype=torch.int32) - bias = torch.zeros(32, device="cuda") - histogram = torch.zeros(32, 31, device="cuda", dtype=torch.int32) - bounds = torch.tensor([-1.0, 1.0], device="cuda") - - def route(x): - return fused_topk_with_score_function( - x, - 4, - False, - None, - None, - None, - "sigmoid", - bias, - topk_indices=indices, - qb_histogram=histogram if histogram_mode else None, - qb_bin_bounds=bounds if histogram_mode else None, - qb_histogram_mode=histogram_mode, - ) - - route = _router_function(route, execution) - probs, _ = route(logits) - route(-logits) - with pytest.raises(RuntimeError, match="modified by an inplace operation"): - probs.sum().backward() + _check_reference(actual_grad, ref_grad, dtype) def test_qb_topk_argument_validation(): From 80b79209975f61a9549ba22aa3cdbd53716f7cc0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 19:34:51 +0200 Subject: [PATCH 10/14] Remove dedicated compiled QB bounds validation tests Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 79 ------------------------------ 1 file changed, 79 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 92779b05ea9..a9cf78badaa 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -333,85 +333,6 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogra ) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -@pytest.mark.parametrize("requires_grad", [False, True]) -def test_qb_compile_skips_value_validation(monkeypatch, histogram_mode, requires_grad): - from transformer_engine.pytorch import router - - torch._dynamo.reset() - logits = torch.randn(8, 16, device="cuda", requires_grad=requires_grad) - bias = torch.zeros(16, device="cuda") - bounds = torch.tensor([-1.0, 1.0], device="cuda") - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - - def route(x): - return fused_topk_with_score_function( - x, - 4, - False, - None, - None, - None, - "sigmoid", - bias, - qb_histogram=histogram, - qb_bin_bounds=bounds, - qb_histogram_mode=histogram_mode, - ) - - def unexpected_validation(_): - raise AssertionError("compiled QB must not read bounds values on the CPU") - - compiled = torch.compile(route, fullgraph=True) - with monkeypatch.context() as patch: - patch.setattr(router, "_validate_qb_bin_bounds", unexpected_validation) - for _ in range(2): - probs, _ = compiled(logits) - if requires_grad: - probs.square().sum().backward() - bounds.add_(0.1) - torch.testing.assert_close(histogram.sum(dim=-1), torch.full((16,), 16, device="cuda")) - bounds.zero_() - with pytest.raises(ValueError, match="finite with lower < upper"): - route(logits) - - -@pytest.mark.parametrize("compiled", [False, True]) -@pytest.mark.parametrize("invalid", ["shape", "dtype", "device"]) -def test_qb_bounds_check_metadata(compiled, invalid): - torch._dynamo.reset() - logits = torch.randn(8, 16, device="cuda") - bias = torch.zeros(16, device="cuda") - histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) - bounds = torch.tensor([-1.0, 1.0], device="cuda") - if invalid == "shape": - bounds = bounds.reshape(1, 2) - elif invalid == "dtype": - bounds = bounds.half() - else: - bounds = bounds.cpu() - - def router(logits, bounds): - return fused_topk_with_score_function( - logits, - 4, - False, - None, - None, - None, - "sigmoid", - bias, - qb_histogram=histogram, - qb_bin_bounds=bounds, - qb_histogram_mode="two_kernel", - ) - - if compiled: - router = torch.compile(router, fullgraph=True) - with pytest.raises(RuntimeError, match="QB bin_bounds"): - router(logits, bounds) - - # Pytorch-based group topk def group_limited_topk( scores: torch.Tensor, From b8d8eecb535c7cc24de6b1e14d6174269b6643ea Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 19:37:53 +0200 Subject: [PATCH 11/14] Document router buffer constraints beside their parameters Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/router.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index b2fcaaf9b37..9122b583fb8 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -357,13 +357,7 @@ def fused_topk_with_score_function( qb_histogram_mode: Optional[str] = None, ): """ - Fused topk with score function router. - - ``torch.compile(fullgraph=True)`` supports routing maps, dense output buffers, and - in-place QB histogram accumulation. Compiled execution skips QB bounds value validation - and its CPU synchronization; tensor metadata is still checked. For automatic CUDA graph - replay, use module buffers or buffers marked with ``torch._dynamo.mark_static_address`` - for mutated inputs. Ordinary mutated inputs may cause CUDA graph capture to be skipped. + Fused topk with score function router. Supports ``torch.compile(fullgraph=True)``. Parameters ---------- @@ -391,12 +385,17 @@ def fused_topk_with_score_function( controls the dense index output dtype and the routing map is not materialized. The buffer is overwritten and returned in both eager and compiled execution. Its contents must remain unchanged until the corresponding backward completes. + For automatic CUDA graph replay, use a module buffer or + ``torch._dynamo.mark_static_address``; otherwise capture may be skipped. qb_histogram : torch.Tensor, optional - Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. + Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place in eager + and compiled execution. For automatic CUDA graph replay, use a module buffer or + ``torch._dynamo.mark_static_address``; otherwise capture may be skipped. qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. - Under ``torch.compile``, the caller must guarantee valid values on every call, + Under ``torch.compile``, value validation and its CPU synchronization are skipped; + tensor metadata is still checked. The caller must guarantee valid values on every call, including after device-side updates; compiled calls do not mark bounds as validated for later eager calls. Before eager CUDA graph capture, validate the current version eagerly or use From 9282a4bd35ed778d506b1fec16d1ab8aeb602f6b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 15 Oct 2026 23:50:53 +0200 Subject: [PATCH 12/14] Use setup_context for compiled fused routing Move context setup out of compiled forward so Dynamo handles grad and no-grad calls without a separate inference predicate in the router. Keep the original-style Function for fast eager execution and share its backward with the compiled Function. Expose the existing intermediate as an internal output and disable unused gradient materialization under compile. The public outputs and buffer mutation contract stay unchanged. Validate with the full router suite (3375 passed, 446 skipped), replay profiles, and 100 interleaved CPU timing samples per variant and case. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/router.py | 256 ++++++++++++++++++--------- 1 file changed, 170 insertions(+), 86 deletions(-) diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 9122b583fb8..73f3c55c783 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -156,41 +156,23 @@ def forward( expert_bias: Optional[torch.Tensor], routing_map_format: int, topk_indices: Optional[torch.Tensor], - skip_context: bool, ): # pylint: disable=missing-function-docstring - compiled = torch.compiler.is_compiling() - if compiled: - probs, routing_output, intermediate_output = _topk_forward( - logits, - topk, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - score_function, - expert_bias, - routing_map_format, - topk_indices, - ) - if skip_context: - return probs, routing_output - else: - probs, routing_output, intermediate_output = tex.fused_topk_with_score_function_fwd( - logits, - topk, - use_pre_softmax, - num_groups, - group_topk, - scaling_factor, - score_function, - expert_bias, - routing_map_format, - topk_indices, - ) - if topk_indices is not None: - routing_output = topk_indices - ctx.compiled = compiled + probs, routing_output, intermediate_output = tex.fused_topk_with_score_function_fwd( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, + ) + if topk_indices is not None: + routing_output = topk_indices + ctx.compiled = False if topk_indices is not None: ctx.mark_dirty(topk_indices) ctx.mark_non_differentiable(routing_output) @@ -204,7 +186,7 @@ def forward( return probs, routing_output @staticmethod - def backward(ctx, grad_probs, _): + def backward(ctx, grad_probs, _, _grad_intermediate=None): # pylint: disable=missing-function-docstring routing_map, intermediate_output = ctx.saved_tensors if ctx.compiled: @@ -219,7 +201,7 @@ def backward(ctx, grad_probs, _): ctx.routing_map_format, ctx.use_dense_indices, ) - return grad_logits, None, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -235,7 +217,7 @@ def backward(ctx, grad_probs, _): ctx.use_dense_indices, ctx.routing_map_format, ) - return grad_logits, None, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None, None class FusedTopkScoreFunctionQB(torch.autograd.Function): @@ -253,46 +235,28 @@ def forward( histogram: torch.Tensor, bin_bounds: torch.Tensor, histogram_mode: int, - skip_context: bool, ): # pylint: disable=missing-function-docstring - compiled = torch.compiler.is_compiling() - if compiled: - probs, routing_output, intermediate_output = _qb_forward( - logits, - topk, - scaling_factor, - expert_bias, - routing_map_format, - topk_indices, - histogram, - bin_bounds, - histogram_mode, - ) - if skip_context: - return probs, routing_output, histogram - histogram_output = histogram - else: - bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) - ( - probs, - routing_output, - intermediate_output, - _cutoff, - histogram_output, - ) = tex.fused_topk_with_score_function_qb_fwd( - logits, - topk, - scaling_factor, - expert_bias, - routing_map_format, - topk_indices, - histogram, - bin_bounds, - histogram_mode, - bin_bounds_validated, - ) - ctx.compiled = compiled + bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) + ( + probs, + routing_output, + intermediate_output, + _cutoff, + histogram_output, + ) = tex.fused_topk_with_score_function_qb_fwd( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + bin_bounds_validated, + ) + ctx.compiled = False if topk_indices is not None: routing_output = topk_indices ctx.mark_dirty(topk_indices, histogram) @@ -307,7 +271,7 @@ def forward( return probs, routing_output, histogram_output @staticmethod - def backward(ctx, grad_probs, _, _histogram_grad): + def backward(ctx, grad_probs, _, _histogram_grad, _grad_intermediate=None): # pylint: disable=missing-function-docstring routing_output, intermediate_output = ctx.saved_tensors if ctx.compiled: @@ -322,7 +286,7 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.routing_map_format, ctx.use_dense_indices, ) - return grad_logits, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None if not grad_probs.is_contiguous(): grad_probs = grad_probs.contiguous() grad_logits = torch.empty_like(grad_probs) @@ -338,7 +302,124 @@ def backward(ctx, grad_probs, _, _histogram_grad): ctx.use_dense_indices, ctx.routing_map_format, ) - return grad_logits, None, None, None, None, None, None, None, None, None + return grad_logits, None, None, None, None, None, None, None, None + + +class _CompiledFusedTopkScoreFunction(FusedTopkScoreFunction): + """ + Fused Topk with Score Function router. + Currently, support "softmax", "sigmoid" and "sqrtsoftplus". + """ + + @staticmethod + def forward( + logits: torch.Tensor, + topk: int, + use_pre_softmax: bool, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, + expert_bias: Optional[torch.Tensor], + routing_map_format: int, + topk_indices: Optional[torch.Tensor], + ): + # pylint: disable=missing-function-docstring + probs, routing_output, intermediate_output = _topk_forward( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + routing_map_format, + topk_indices, + ) + return probs, routing_output, intermediate_output + + @staticmethod + def setup_context(ctx, inputs, output): + # pylint: disable=missing-function-docstring + ( + _, + ctx.topk, + ctx.use_pre_softmax, + _, + _, + ctx.scaling_factor, + ctx.score_function, + _, + ctx.routing_map_format, + topk_indices, + ) = inputs + _, routing_output, intermediate_output = output + ctx.compiled = True + if topk_indices is not None: + ctx.mark_dirty(topk_indices) + ctx.mark_non_differentiable(routing_output, intermediate_output) + ctx.set_materialize_grads(False) + ctx.save_for_backward(routing_output, intermediate_output) + ctx.use_dense_indices = topk_indices is not None + + +class _CompiledFusedTopkScoreFunctionQB(FusedTopkScoreFunctionQB): + """Kimi K3 QB router with histogram accumulation.""" + + @staticmethod + def forward( + logits: torch.Tensor, + topk: int, + scaling_factor: Optional[float], + expert_bias: torch.Tensor, + routing_map_format: int, + topk_indices: Optional[torch.Tensor], + histogram: torch.Tensor, + bin_bounds: torch.Tensor, + histogram_mode: int, + ): + # pylint: disable=missing-function-docstring + probs, routing_output, intermediate_output = _qb_forward( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + ) + histogram_output = histogram + if topk_indices is not None: + routing_output = topk_indices + return probs, routing_output, histogram_output, intermediate_output + + @staticmethod + def setup_context(ctx, inputs, output): + # pylint: disable=missing-function-docstring + ( + _, + ctx.topk, + ctx.scaling_factor, + _, + ctx.routing_map_format, + topk_indices, + histogram, + _, + _, + ) = inputs + _, routing_output, histogram_output, intermediate_output = output + ctx.compiled = True + if topk_indices is not None: + ctx.mark_dirty(topk_indices, histogram) + else: + ctx.mark_dirty(histogram) + ctx.mark_non_differentiable(routing_output, histogram_output, intermediate_output) + ctx.set_materialize_grads(False) + ctx.save_for_backward(routing_output, intermediate_output) + ctx.use_dense_indices = topk_indices is not None def fused_topk_with_score_function( @@ -423,10 +504,6 @@ def fused_topk_with_score_function( if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type.") routing_map_format = _validate_routing_map_format(routing_map_format) - compiled_inference = torch.compiler.is_compiling() and ( - not torch.is_grad_enabled() - or not any(t is not None and t.requires_grad for t in (logits, expert_bias, qb_bin_bounds)) - ) qb_arguments = (qb_histogram, qb_bin_bounds, qb_histogram_mode) if any(value is not None for value in qb_arguments): if any(value is None for value in qb_arguments): @@ -445,7 +522,12 @@ def fused_topk_with_score_function( "qb_histogram_mode must be 'two_kernel' or 'fused_atomic', " f"got {qb_histogram_mode!r}" ) - probs, routing_output, _ = FusedTopkScoreFunctionQB.apply( + function = ( + _CompiledFusedTopkScoreFunctionQB + if torch.compiler.is_compiling() + else FusedTopkScoreFunctionQB + ) + outputs = function.apply( logits, topk, scaling_factor, @@ -455,10 +537,12 @@ def fused_topk_with_score_function( qb_histogram, qb_bin_bounds, mode, - compiled_inference, ) - return probs, routing_output - return FusedTopkScoreFunction.apply( + return outputs[:2] + function = ( + _CompiledFusedTopkScoreFunction if torch.compiler.is_compiling() else FusedTopkScoreFunction + ) + outputs = function.apply( logits, topk, use_pre_softmax, @@ -469,8 +553,8 @@ def fused_topk_with_score_function( expert_bias, routing_map_format, topk_indices, - compiled_inference, ) + return outputs[:2] class FusedComputeScoresForMoEAuxLoss(torch.autograd.Function): From eb9e922c996c3eaaaaf329457643b0d45e6b9bff Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 16 Oct 2026 12:24:32 +0200 Subject: [PATCH 13/14] Check that dynamic router inputs do not trigger recompilation Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index a9cf78badaa..d6d8c83332f 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -226,6 +226,7 @@ def router(logits, bounds, total, indices, histogram): ) def test_router_compile_dynamic(dtype, index_dtype): torch._dynamo.reset() + counters.clear() counts = torch.arange(1, 131, device="cuda", dtype=torch.int32) def router_loss(logits): @@ -262,6 +263,7 @@ def router_loss(logits): (actual_grad,) = torch.autograd.grad(actual, logits) (eager_grad,) = torch.autograd.grad(expected, eager_logits) _check_reference(actual_grad, eager_grad, dtype) + assert counters["stats"]["unique_graphs"] == 1 @pytest.mark.parametrize( From c2fb139df19e44f45272de9398f0204e547dd3da Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 16 Oct 2026 12:32:05 +0200 Subject: [PATCH 14/14] Check fused router backward custom operators with opcheck Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fused_router.py | 94 +++++++++++++++++++----------- 1 file changed, 60 insertions(+), 34 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index d6d8c83332f..7bc42b3fb42 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -295,44 +295,70 @@ def test_router_compile_opcheck(dtype, routing_map_format, index_dtype, histogra bias = torch.zeros(130, device="cuda") histogram = torch.full((130, 64), 7, device="cuda", dtype=torch.int32) bounds = torch.tensor([-1.0, 1.0], device="cuda") - torch.library.opcheck( - router._qb_indices_forward if indices is not None else router._qb_map_forward, - ( - forward_logits, - 4, - None, - bias, - routing_map_format, - *((indices,) if indices is not None else ()), - histogram, - bounds, - int(histogram_mode), - ), + forward = router._qb_indices_forward if indices is not None else router._qb_map_forward + forward_args = ( + forward_logits, + 4, + None, + bias, + routing_map_format, + *((indices,) if indices is not None else ()), + histogram, + bounds, + int(histogram_mode), ) else: + forward = router._topk_indices_forward if indices is not None else router._topk_map_forward + forward_args = ( + forward_logits, + 4, + False, + None, + None, + None, + "sigmoid", + None, + routing_map_format, + *((indices,) if indices is not None else ()), + ) + torch.library.opcheck(forward, forward_args) + outputs = forward(*forward_args) + torch.library.opcheck( + router._topk_backward, + ( + indices if indices is not None else outputs[1], + outputs[-1], + torch.randn_like(outputs[0]), + 4, + False, + None, + "sigmoid", + routing_map_format, + indices is not None, + ), + ) + + if histogram_mode is None and index_dtype in (None, torch.int16): + torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) + scores, _, intermediate = router._score_forward( + forward_logits, 4, "softmax", routing_map_format + ) torch.library.opcheck( - router._topk_indices_forward if indices is not None else router._topk_map_forward, - ( - forward_logits, - 4, - False, - None, - None, - None, - "sigmoid", - None, - routing_map_format, - *((indices,) if indices is not None else ()), - ), + router._score_backward, + (intermediate, torch.randn_like(scores), 4, "softmax", dtype), ) - if index_dtype in (None, torch.int16): - torch.library.opcheck(router._score_forward, (logits, 4, "softmax", routing_map_format)) - counts = torch.ones(130, device="cuda", dtype=torch.int32) - count = torch.tensor(16, device="cuda", dtype=torch.int64) - for token_count in (None, count): - torch.library.opcheck( - router._aux_loss_forward, (logits, counts, 16, token_count, 130, 4, 0.01) - ) + counts = torch.ones(130, device="cuda", dtype=torch.int32) + count = torch.tensor(16, device="cuda", dtype=torch.int64) + for token_count in (None, count): + torch.library.opcheck( + router._aux_loss_forward, (logits, counts, 16, token_count, 130, 4, 0.01) + ) + loss, const_buf = router._aux_loss_forward( + forward_logits, counts, 16, token_count, 130, 4, 0.01 + ) + torch.library.opcheck( + router._aux_loss_backward, (const_buf, counts, 16, 130, torch.randn_like(loss)) + ) # Pytorch-based group topk