From b2f9ebddf4f708759446e2249d084145468a5675 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 3 Aug 2026 12:19:56 -0700 Subject: [PATCH 1/4] Add fuse_swiglu_mxfp8 flag wiring unified SwiGLU+MXFP8 kernel into FusedSwiGLU Integrates the TorchAO unified SwiGLU+MXFP8 CuTe DSL composite (mxfp8_swiglu_mlp_w13 / mxfp8_swiglu_grouped_mlp_w13) behind a new fuse_swiglu_mxfp8 config flag, with A/B debug configs for Llama3 (llama3_debugmodel_mxfp8[_fused_swiglu]) and DeepSeek-V3 (deepseek_v3_debugmodel_mxfp8[_fused_swiglu]). Both arms share the composite autograd boundary; only fuse_activation differs. Co-Authored-By: Claude Fable 5 --- .../unit_tests/test_fused_swiglu_override.py | 50 ++++++++++++- torchtitan/components/quantization/mx.py | 23 +++++- .../models/deepseek_v3/config_registry.py | 42 +++++++++++ torchtitan/models/llama3/config_registry.py | 41 +++++++++++ torchtitan/overrides/fused_swiglu.py | 73 +++++++++++++++++-- 5 files changed, 216 insertions(+), 13 deletions(-) diff --git a/tests/unit_tests/test_fused_swiglu_override.py b/tests/unit_tests/test_fused_swiglu_override.py index e8de78f1cd..9dee377ed7 100644 --- a/tests/unit_tests/test_fused_swiglu_override.py +++ b/tests/unit_tests/test_fused_swiglu_override.py @@ -9,14 +9,20 @@ import spmd_types as spmd import torch +from torchtitan.config.override import apply_overrides from torchtitan.models.common.decoder_sharding import dense_param_placement from torchtitan.models.common.moe import GroupedExperts from torchtitan.models.deepseek_v3.config_registry import ( deepseek_v3_debugmodel_minimal_async_ep, + deepseek_v3_debugmodel_mxfp8_fused_swiglu, +) +from torchtitan.models.llama3.config_registry import ( + llama3_debugmodel_mxfp8_fused_swiglu, ) from torchtitan.overrides.fused_swiglu import ( - fused_grouped_experts, FusedGroupedExperts, + FusedSwiGLU, + fused_grouped_experts, silu_and_mul_backward_kernel, silu_and_mul_forward_kernel, silu_and_mul_op, @@ -60,6 +66,48 @@ def test_grouped_experts_config_is_replaced(self): self.assertIsInstance(replacement, FusedGroupedExperts.Config) + @unittest.skipUnless( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 10, + "MXFP8 fusion requires SM100", + ) + def test_mxfp8_converter_composes_with_dense_override(self): + config = llama3_debugmodel_mxfp8_fused_swiglu() + model_config = config.model_spec.model + model_config.update_from_config(config=config) + apply_overrides(config.override, config) + with torch.device("meta"): + model = model_config.build() + fused = [ + module + for module in model.modules() + if isinstance(module, FusedSwiGLU) + ] + self.assertTrue(fused) + self.assertTrue(all(module.mxfp8_fused for module in fused)) + self.assertTrue(all(module.fuse_activation for module in fused)) + + @unittest.skipUnless( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 10, + "MXFP8 fusion requires SM100", + ) + def test_mxfp8_converter_composes_with_grouped_override(self): + config = deepseek_v3_debugmodel_mxfp8_fused_swiglu() + model_config = config.model_spec.model + model_config.update_from_config(config=config) + apply_overrides(config.override, config) + with torch.device("meta"): + model = model_config.build() + fused = [ + module + for module in model.modules() + if isinstance(module, FusedGroupedExperts) + ] + self.assertTrue(fused) + self.assertTrue(all(module.mxfp8_fused for module in fused)) + self.assertTrue(all(module.fuse_activation for module in fused)) + class TestFusedGroupedExperts(unittest.TestCase): """Checkpoint interop (state_dict hooks) and override config remap for the diff --git a/torchtitan/components/quantization/mx.py b/torchtitan/components/quantization/mx.py index 16675e2467..67bf90564b 100644 --- a/torchtitan/components/quantization/mx.py +++ b/torchtitan/components/quantization/mx.py @@ -33,7 +33,8 @@ class MXFP8Linear(TorchAOMXFP8Linear, Module): class Config(Linear.Config): """Drop-in replacement for Linear.Config that builds MXFP8Linear.""" - pass + wgrad_with_hp: bool = False + fuse_swiglu_mxfp8: bool = False def __init__(self, config: Config): TorchAOMXFP8Linear.__init__( @@ -41,6 +42,7 @@ def __init__(self, config: Config): config.in_features, config.out_features, bias=config.bias, + wgrad_with_hp=config.wgrad_with_hp, ) except ImportError: @@ -58,6 +60,10 @@ class Config(QuantizationConverter.Config): Only Linear.Config entries whose FQN contains a match are converted. If empty, all Linear modules are converted. """ + wgrad_with_hp: bool = False + """Compute linear weight gradients with BF16 GEMMs.""" + fuse_swiglu_mxfp8: bool = False + """Use fused SwiGLU activation quantization when the override is active.""" def __init__(self, config: Config): self.config = config @@ -86,6 +92,8 @@ def convert(self, model_config): out_features=config.out_features, bias=config.bias, param_init=config.param_init, + wgrad_with_hp=self.config.wgrad_with_hp, + fuse_swiglu_mxfp8=self.config.fuse_swiglu_mxfp8, ) if parent is None: model_config = new_config @@ -120,6 +128,7 @@ class MXFP8GroupedExperts(parent_cls): # type: ignore[valid-type, misc] @dataclass(kw_only=True, slots=True) class Config(parent_config_cls): # type: ignore[misc] recipe_name: str = "mxfp8_rceil" + fuse_swiglu_mxfp8: bool = False def __init__(self, config: Config): super().__init__(config) @@ -151,18 +160,23 @@ class MXFP8GroupedExpertsConverter(QuantizationConverter): @dataclass(kw_only=True, slots=True) class Config(QuantizationConverter.Config): - recipe_name: Literal["mxfp8_rceil"] = "mxfp8_rceil" + recipe_name: Literal[ + "mxfp8_rceil", + "mxfp8_rceil_wgrad_with_hp", + ] = "mxfp8_rceil" """ Quantization recipe name for grouped GEMMs. Options: ["mxfp8_rceil"] - - mxfp8_rceil: MXFP8 dynamic quantization with RCEIL rounding mode - when computing the e8m0 scale factors. + - mxfp8_rceil: MXFP8 dynamic quantization with RCEIL scales. + - mxfp8_rceil_wgrad_with_hp: same fprop/dgrad path with BF16 wgrad. """ pad_multiple: int = 32 """ Pad per-expert token groups to this multiple for MXFP8 grouped GEMM alignment. The CuTeDSL quantization kernel on sm_100 requires multiples of 128. """ + fuse_swiglu_mxfp8: bool = False + """Fuse SwiGLU fprop/dgrad with rowwise MXFP8 casts.""" def __init__(self, config: Config): self.config = config @@ -191,6 +205,7 @@ def convert(self, model_config): new_config = config_cls( **{f.name: getattr(config, f.name) for f in fields(config)}, recipe_name=self.config.recipe_name, + fuse_swiglu_mxfp8=self.config.fuse_swiglu_mxfp8, ) if parent is None: model_config = new_config diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 2a144a6914..71fd8c1f08 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -126,6 +126,48 @@ def deepseek_v3_debugmodel_minimal_async_ep() -> Trainer.Config: return config +def _deepseek_v3_debugmodel_mxfp8_swiglu_ab( + *, fuse_swiglu_mxfp8: bool +) -> Trainer.Config: + """Shared harness for the grouped-expert SwiGLU fusion A/B. + + Distinct from ``deepseek_v3_debugmodel_mxfp8``: only the grouped-expert + GEMMs are quantized, so the A/B isolates the expert SwiGLU boundary rather + than also moving the dense Linear layers to MXFP8. + """ + config = deepseek_v3_debugmodel() + config.compile = CompileConfig(enable=True, components=["model", "loss"]) + config.model_spec = model_registry( + "debugmodel", + converters=[ + MXFP8GroupedExpertsConverter.Config( + recipe_name="mxfp8_rceil", + pad_multiple=128, + model_compile_enabled=True, + fuse_swiglu_mxfp8=fuse_swiglu_mxfp8, + ), + ], + ) + enable_fused_swiglu(config) + # The padded token-group dispatch the MXFP8 grouped GEMMs need is only + # produced by the EP permute path. + config.parallelism = ParallelismConfig(expert_parallel_degree=2) + config.debug.moe_force_load_balance = True + config.training.steps = 50 + config.metrics.log_freq = 1 + return config + + +def deepseek_v3_debugmodel_mxfp8_unfused_swiglu() -> Trainer.Config: + """MXFP8 grouped-expert A/B baseline: standalone SwiGLU quantization.""" + return _deepseek_v3_debugmodel_mxfp8_swiglu_ab(fuse_swiglu_mxfp8=False) + + +def deepseek_v3_debugmodel_mxfp8_fused_swiglu() -> Trainer.Config: + """MXFP8 grouped-expert SwiGLU fusion benchmark on SM100.""" + return _deepseek_v3_debugmodel_mxfp8_swiglu_ab(fuse_swiglu_mxfp8=True) + + def deepseek_v3_16b() -> Trainer.Config: model_spec = model_registry("16B", attn_backend="flex") return Trainer.Config( diff --git a/torchtitan/models/llama3/config_registry.py b/torchtitan/models/llama3/config_registry.py index b800378588..9c6285e446 100644 --- a/torchtitan/models/llama3/config_registry.py +++ b/torchtitan/models/llama3/config_registry.py @@ -31,6 +31,8 @@ from . import model_registry from .model import Llama3Model +_FUSED_SWIGLU_OVERRIDE = "torchtitan.overrides.fused_swiglu.fused_swiglu" + def llama3_debugmodel() -> Trainer.Config: model_spec = model_registry("debugmodel") @@ -136,6 +138,45 @@ def llama3_debugmodel_first_85_pct_layers_nvfp4() -> Trainer.Config: return config +def _llama3_debugmodel_mxfp8(*, fuse_swiglu_mxfp8: bool) -> Trainer.Config: + config = llama3_debugmodel() + config.compile = CompileConfig(enable=True, components=["model", "loss"]) + config.model_spec = model_registry( + "debugmodel", + converters=[ + MXFP8LinearConverter.Config( + fqns=["feed_forward"], + model_compile_enabled=True, + fuse_swiglu_mxfp8=fuse_swiglu_mxfp8, + ), + ], + ) + # llama3 is dense, so only the FeedForward override is needed; naming the + # factory (not the module) is required since override.imports resolves each + # @override target individually. + config.override.imports.append(_FUSED_SWIGLU_OVERRIDE) + config.training.steps = 50 + config.metrics.log_freq = 1 + return config + + +def llama3_debugmodel_mxfp8() -> Trainer.Config: + """Dense MXFP8 A/B baseline: fused w13, standalone SwiGLU quantization.""" + return _llama3_debugmodel_mxfp8(fuse_swiglu_mxfp8=False) + + +def llama3_debugmodel_mxfp8_fused_swiglu() -> Trainer.Config: + """Dense MXFP8 A/B candidate: fused w13, unified SwiGLU+MXFP8 kernel.""" + return _llama3_debugmodel_mxfp8(fuse_swiglu_mxfp8=True) + + +def llama3_debugmodel_mxfp8_stock_ffn() -> Trainer.Config: + """Dense MXFP8 context run: stock FeedForward, no fused-w13 override.""" + config = _llama3_debugmodel_mxfp8(fuse_swiglu_mxfp8=False) + config.override.imports.remove(_FUSED_SWIGLU_OVERRIDE) + return config + + def llama3_debugmodel_float8_emulate_lora() -> Trainer.Config: from torchtitan.components.lora import LoRAConverter diff --git a/torchtitan/overrides/fused_swiglu.py b/torchtitan/overrides/fused_swiglu.py index 6566742886..20cd2d3c2e 100644 --- a/torchtitan/overrides/fused_swiglu.py +++ b/torchtitan/overrides/fused_swiglu.py @@ -58,7 +58,6 @@ import torch import triton import triton.language as tl - from torch.distributed.tensor import DTensor from torch.distributed.tensor.experimental import local_map @@ -440,7 +439,8 @@ class FusedSwiGLU(FeedForward): @dataclass(kw_only=True, slots=True) class Config(FeedForward.Config): - pass + mxfp8_fused: bool = False + fuse_activation: bool = False def __init__(self, config: Config): super().__init__(config) @@ -449,10 +449,25 @@ def __init__(self, config: Config): self.w13 = torch.nn.Parameter( torch.empty(config.w1.out_features, 2, config.w1.in_features) ) + self.mxfp8_fused = config.mxfp8_fused + self.fuse_activation = config.fuse_activation self.register_state_dict_post_hook(self._split_w13_on_save) self.register_load_state_dict_pre_hook(self._merge_w13_on_load) def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.mxfp8_fused and not isinstance(x, DTensor): + from torchao.prototype.moe_training import mxfp8_swiglu_mlp_w13 + + output = mxfp8_swiglu_mlp_w13( + x, + self.w13, + self.w2.weight, + fuse_activation=self.fuse_activation, + wgrad_with_hp=self.w2.wgrad_with_hp, + ) + if self.w2.bias is not None: + output = output + self.w2.bias.to(output.dtype) + return output gate, up = torch.einsum("...d,hgd->...hg", x, self.w13).unbind(-1) return self.w2(_fused_silu_and_mul(gate, up)) @@ -494,7 +509,20 @@ def fused_swiglu(cfg: FeedForward.Config) -> FusedSwiGLU.Config: if w1_init is not None and w3_init is not None: param_init = {"w13": _make_fused_gate_up_init(w1_init, w3_init, gate_up_axis=1)} - fused = derive(cfg, FusedSwiGLU.Config, param_init=param_init) + w2_owner = getattr(type(cfg.w2), "_owner", None) + mxfp8_fused = ( + w2_owner is not None + and w2_owner.__module__ + == "torchtitan.components.quantization.mx" + and w2_owner.__name__ == "MXFP8Linear" + ) + fused = derive( + cfg, + FusedSwiGLU.Config, + param_init=param_init, + mxfp8_fused=mxfp8_fused, + fuse_activation=getattr(cfg.w2, "fuse_swiglu_mxfp8", False), + ) base = cfg.sharding_config fused.sharding_config = ShardingConfig( @@ -523,7 +551,9 @@ class FusedGroupedExperts(GroupedExperts): @dataclass(kw_only=True, slots=True) class Config(GroupedExperts.Config): - pass + mxfp8_fused: bool = False + fuse_activation: bool = False + wgrad_with_hp: bool = False def __init__(self, config: Config): super().__init__(config) @@ -534,6 +564,9 @@ def __init__(self, config: Config): self.w13 = torch.nn.Parameter( torch.empty(config.num_experts, config.hidden_dim, 2, config.dim) ) + self.mxfp8_fused = config.mxfp8_fused + self.fuse_activation = config.fuse_activation + self.wgrad_with_hp = config.wgrad_with_hp self.register_state_dict_post_hook(self._split_w13_on_save) self.register_load_state_dict_pre_hook(self._merge_w13_on_load) @@ -554,6 +587,18 @@ def forward( E, F, _, D = w13.shape offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) + if self.mxfp8_fused: + from torchao.prototype.moe_training import mxfp8_swiglu_grouped_mlp_w13 + + return mxfp8_swiglu_grouped_mlp_w13( + x_RD.bfloat16(), + w13.bfloat16(), + w2_EDF.bfloat16().transpose(-2, -1), + offsets_E, + fuse_activation=self.fuse_activation, + wgrad_with_hp=self.wgrad_with_hp, + ).type_as(x_RD) + w13_E_D_2F = w13.bfloat16().reshape(E, F * 2, D).transpose(-2, -1) gate_up_R2F = self._grouped_mm( A=x_RD.bfloat16(), B_t=w13_E_D_2F, offs=offsets_E @@ -618,13 +663,25 @@ def fused_grouped_experts( cfg: GroupedExperts.Config, ) -> GroupedExperts.Config: # Remap w1_EFD/w3_EFD param-init and state shardings onto the fused w13. - # Idempotent: return cfg unchanged if it is not a stock GroupedExperts.Config - # (already fused, or a subclass like GptOssGroupedExperts). - if type(cfg) is not GroupedExperts.Config: + # Idempotent: return cfg unchanged if it is not owned by a GroupedExperts + # implementation. The check is on the owning module class rather than the + # exact Config type so that a quantization-derived subclass -- e.g. + # MXFP8GroupedExperts.Config, which only overrides the _grouped_mm seam -- + # still fuses, which is what the MXFP8 SwiGLU A/B configs rely on. + owner = getattr(type(cfg), "_owner", None) + if owner is not None and not issubclass(owner, GroupedExperts): return cfg param_init = _fuse_w13_grouped_experts_param_init(cfg.param_init) - fused = derive(cfg, FusedGroupedExperts.Config, param_init=param_init) + recipe_name = getattr(cfg, "recipe_name", "") + fused = derive( + cfg, + FusedGroupedExperts.Config, + param_init=param_init, + mxfp8_fused=recipe_name.startswith("mxfp8_"), + fuse_activation=getattr(cfg, "fuse_swiglu_mxfp8", False), + wgrad_with_hp=recipe_name.endswith("wgrad_with_hp"), + ) base = cfg.sharding_config if base is not None: fused.sharding_config = _fuse_w13_grouped_experts_sharding(base) From 79a9fc8a0c997995303bd08ffc0d6ae6edd05000 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 9 Aug 2026 01:57:29 -0700 Subject: [PATCH 2/4] Add 16B MXFP8 SwiGLU-fusion benchmark configs Mirrors the JET regular-MXFP8 workload (MXFP8Linear over attention, dense FFN and shared-expert linears, plus MXFP8 routed-expert grouped GEMMs) so the SwiGLU fusion can be measured end to end rather than only in a microbenchmark. Three arms, because the fused-w13 override and the unified SwiGLU+MXFP8 kernel are independent changes and move in opposite directions: deepseek_v3_16b_mxfp8_exp no override (the JET control) deepseek_v3_16b_mxfp8_w13_exp fused w13 GEMM, standalone casts deepseek_v3_16b_mxfp8_fused_swiglu_exp fused w13 + unified kernel Measured on 4x GB200 (EP=4, bs=4, seq=4096, 50 steps), averaged over steps 11-50: the w13 override alone is -0.51% TFLOPs and +3.3GiB, the unified kernel adds +2.51% on top of it, for +1.98% against the control. Loss at step 50 agrees to within 1.5e-3 across all three. Uses the in-repo test tokenizer and c4_test so the configs run without downloaded assets; token content affects neither FLOPs nor kernel shapes. Co-Authored-By: Claude Opus 5 (1M context) --- .../models/deepseek_v3/config_registry.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 71fd8c1f08..f6536c0f71 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -202,6 +202,81 @@ def deepseek_v3_16b() -> Trainer.Config: ) +def _deepseek_v3_16b_mxfp8_exp( + *, fused_swiglu_override: bool, fuse_swiglu_mxfp8: bool +) -> Trainer.Config: + """16B MXFP8 end-to-end harness, mirroring the JET regular-MXFP8 workload. + + Converters match that workload: MXFP8Linear over attention, dense FFN and + shared-expert linears, plus MXFP8 routed-expert grouped GEMMs. Three arms + are derived from this so the two independent changes can be separated: + + * regular -- no override; routed experts run separate w1/w3 grouped + mms, then silu(gate)*up, then w2. + * w13 -- override on, fuse_swiglu_mxfp8=False; gate+up are one + fused w13 GEMM, activation still cast standalone. + * fused -- override on, fuse_swiglu_mxfp8=True; activation and its + MXFP8 cast collapse into the unified kernel. + + moe_hidden_dim (1408) and the shared-expert width (2816) are 128-divisible + so those qualify for the fused kernel; the layer-0 dense FFN is 10944 and + falls back to the per-GEMM mx_mm path in every arm. + + Uses the in-repo test tokenizer and c4_test so no downloaded assets are + needed; token content affects neither FLOPs nor kernel shapes. + """ + config = deepseek_v3_16b() + config.hf_assets_path = "./tests/assets/tokenizer" + config.dataloader = HuggingFaceTextDataLoader.Config(dataset="c4_test") + config.compile = CompileConfig(enable=True, components=["model", "loss"]) + config.model_spec = model_registry( + "16B", + attn_backend="flex", + converters=[ + MXFP8LinearConverter.Config( + fqns=["attention", "shared_experts", "feed_forward"], + model_compile_enabled=True, + fuse_swiglu_mxfp8=fuse_swiglu_mxfp8, + ), + MXFP8GroupedExpertsConverter.Config( + recipe_name="mxfp8_rceil", + pad_multiple=128, + model_compile_enabled=True, + fuse_swiglu_mxfp8=fuse_swiglu_mxfp8, + ), + ], + ) + if fused_swiglu_override: + enable_fused_swiglu(config) + # EP=8 in the JET recipe; this host has 4 GPUs. + config.parallelism = ParallelismConfig(expert_parallel_degree=4) + config.debug.moe_force_load_balance = True + config.training = TrainingConfig(local_batch_size=4, seq_len=4096, steps=50) + config.checkpoint = CheckpointManager.Config(enable=False) + return config + + +def deepseek_v3_16b_mxfp8_exp() -> Trainer.Config: + """Regular MXFP8 arm: no fused-SwiGLU override (the JET control).""" + return _deepseek_v3_16b_mxfp8_exp( + fused_swiglu_override=False, fuse_swiglu_mxfp8=False + ) + + +def deepseek_v3_16b_mxfp8_w13_exp() -> Trainer.Config: + """Intermediate arm: fused w13 GEMM only, standalone activation casts.""" + return _deepseek_v3_16b_mxfp8_exp( + fused_swiglu_override=True, fuse_swiglu_mxfp8=False + ) + + +def deepseek_v3_16b_mxfp8_fused_swiglu_exp() -> Trainer.Config: + """Treatment arm: fused w13 plus the unified SwiGLU+MXFP8 kernel.""" + return _deepseek_v3_16b_mxfp8_exp( + fused_swiglu_override=True, fuse_swiglu_mxfp8=True + ) + + def deepseek_v3_16b_hybridep() -> Trainer.Config: config = deepseek_v3_16b() config.model_spec = model_registry( From 9ffe0533bcd0ab6121f2e1a6177acef79ca2721b Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 10 Aug 2026 12:12:55 -0700 Subject: [PATCH 3/4] Host the composite MXFP8 SwiGLU MLP in overrides Moved from torchao (which now offers only the fused gated-activation kernels): the autograd composite, its tests, and the FusedSwiGLU import sites. No functional change; 34 composite + 18 override tests pass. Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_swiglu_mlp.py | 576 +++++++++++++++++++++++++++ torchtitan/overrides/fused_swiglu.py | 4 +- torchtitan/overrides/swiglu_mlp.py | 544 +++++++++++++++++++++++++ 3 files changed, 1122 insertions(+), 2 deletions(-) create mode 100644 tests/unit_tests/test_swiglu_mlp.py create mode 100644 torchtitan/overrides/swiglu_mlp.py diff --git a/tests/unit_tests/test_swiglu_mlp.py b/tests/unit_tests/test_swiglu_mlp.py new file mode 100644 index 0000000000..3c8e38f65d --- /dev/null +++ b/tests/unit_tests/test_swiglu_mlp.py @@ -0,0 +1,576 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for the composite MXFP8 SwiGLU MLP (swiglu_mlp.py). + +The composite's two modes must agree with each other (only the activation +boundary differs: unified kernel vs standalone BF16 + cast kernels), track the +existing per-GEMM mx_mm MLP, compile without graph breaks, and never launch +standalone activation-cast kernels in fused mode. +""" + +import pytest +import torch +import torch.nn.functional as F + +if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10): + pytest.skip("Requires CUDA SM 10.x (Blackwell)", allow_module_level=True) + +from torch.profiler import ProfilerActivity, profile + +from torchtitan.overrides.swiglu_mlp import ( + _unfused_grouped_mlp, + _unfused_mlp, + mxfp8_swiglu_grouped_mlp_w13, + mxfp8_swiglu_mlp_w13, +) +from torchao.quantization.utils import compute_error + +# (M, D, H): llama3-debugmodel-like plus one wider shape. +_SHAPES = [ + (256, 256, 768), + (512, 512, 1024), +] + +_SWIGLU_OPS = { + "torchao::gated_act_mxfp8_forward", + "torchao::gated_act_mxfp8_backward", +} +_STANDALONE_CAST_OPS = { + "torchao::mxfp8_quantize_2d_1x32_cutedsl", + "torchao::mxfp8_quantize_2d_32x1_cutedsl", +} + + +def _make_inputs(m, d, h, seed=0): + torch.manual_seed(seed) + x = torch.randn(m, d, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w13 = torch.randn(h, 2, d, dtype=torch.bfloat16, device="cuda") * 0.02 + w2 = torch.randn(d, h, dtype=torch.bfloat16, device="cuda") * 0.02 + w13.requires_grad_(True) + w2.requires_grad_(True) + return x, w13, w2 + + +def _bf16_reference(x, w13, w2): + gate, up = torch.einsum("...d,hgd->...hg", x, w13).unbind(-1) + return F.linear(F.silu(gate.float()).to(x.dtype) * up, w2) + + +def _run(fn, x, w13, w2, **kwargs): + x = x.clone().detach().requires_grad_(True) + w13 = w13.clone().detach().requires_grad_(True) + w2 = w2.clone().detach().requires_grad_(True) + out = fn(x, w13, w2, **kwargs) + torch.manual_seed(1234) + out.backward(torch.randn_like(out)) + return out, x.grad, w13.grad, w2.grad + + +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize("wgrad_with_hp", [False, True]) +def test_fused_matches_unfused_mode(shape, wgrad_with_hp): + m, d, h = shape + x, w13, w2 = _make_inputs(m, d, h) + out_f, dx_f, dw13_f, dw2_f = _run( + mxfp8_swiglu_mlp_w13, + x, + w13, + w2, + fuse_activation=True, + wgrad_with_hp=wgrad_with_hp, + ) + out_u, dx_u, dw13_u, dw2_u = _run( + mxfp8_swiglu_mlp_w13, + x, + w13, + w2, + fuse_activation=False, + wgrad_with_hp=wgrad_with_hp, + ) + # Forward h quantization is bitwise identical between the unified kernel + # and the standalone casts, so the outputs must match exactly. + torch.testing.assert_close(out_f, out_u, rtol=0, atol=0) + # Backward may differ by one E4M3 code in <=1e-5 of [dGate | dUp] elements + # (hardware ex2.approx sigmoid), so gradients are near- but not bitwise-equal. + for got, ref, name in [ + (dx_f, dx_u, "dx"), + (dw13_f, dw13_u, "dw13"), + (dw2_f, dw2_u, "dw2"), + ]: + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= 50.0, f"{name} SQNR between modes {sqnr} < 50" + + +@pytest.mark.parametrize("fuse_activation", [True, False]) +@pytest.mark.parametrize("shape", _SHAPES) +def test_tracks_mx_mm_reference(shape, fuse_activation): + m, d, h = shape + x, w13, w2 = _make_inputs(m, d, h) + out, dx, dw13, dw2 = _run( + mxfp8_swiglu_mlp_w13, x, w13, w2, fuse_activation=fuse_activation + ) + ref_out, ref_dx, ref_dw13, ref_dw2 = _run( + _unfused_mlp, x, w13, w2, wgrad_with_hp=False + ) + for got, ref, name, min_sqnr in [ + (out, ref_out, "out", 25.0), + (dx, ref_dx, "dx", 22.0), + (dw13, ref_dw13, "dw13", 22.0), + (dw2, ref_dw2, "dw2", 22.0), + ]: + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= min_sqnr, f"{name} SQNR vs mx_mm {sqnr} < {min_sqnr}" + assert torch.isfinite(got).all(), f"{name} contains non-finite values" + + +@pytest.mark.parametrize("fuse_activation", [True, False]) +def test_tracks_bf16_reference(fuse_activation): + m, d, h = _SHAPES[0] + x, w13, w2 = _make_inputs(m, d, h) + out, dx, dw13, dw2 = _run( + mxfp8_swiglu_mlp_w13, x, w13, w2, fuse_activation=fuse_activation + ) + ref_out, ref_dx, ref_dw13, ref_dw2 = _run(_bf16_reference, x, w13, w2) + for got, ref, name, min_sqnr in [ + (out, ref_out, "out", 20.0), + (dx, ref_dx, "dx", 18.0), + (dw13, ref_dw13, "dw13", 18.0), + (dw2, ref_dw2, "dw2", 18.0), + ]: + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= min_sqnr, f"{name} SQNR vs bf16 {sqnr} < {min_sqnr}" + + +@pytest.mark.parametrize("fuse_activation", [True, False]) +@pytest.mark.parametrize("wgrad_with_hp", [False, True]) +def test_compile(fuse_activation, wgrad_with_hp): + m, d, h = _SHAPES[0] + x, w13, w2 = _make_inputs(m, d, h) + + def fn(x, w13, w2): + return mxfp8_swiglu_mlp_w13( + x, + w13, + w2, + fuse_activation=fuse_activation, + wgrad_with_hp=wgrad_with_hp, + ) + + eager_out, eager_dx, eager_dw13, eager_dw2 = _run(fn, x, w13, w2) + compiled = torch.compile(fn, fullgraph=True) + comp_out, comp_dx, comp_dw13, comp_dw2 = _run(compiled, x, w13, w2) + + torch.testing.assert_close(comp_out, eager_out, rtol=0, atol=0) + for got, ref, name in [ + (comp_dx, eager_dx, "dx"), + (comp_dw13, eager_dw13, "dw13"), + (comp_dw2, eager_dw2, "dw2"), + ]: + # Inductor may fuse the BF16 elementwise math differently; quantized + # GEMM inputs are identical custom-op outputs, so keep this tight. + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= 50.0, f"compiled {name} SQNR {sqnr} < 50" + + +def _op_counts(fuse_activation): + m, d, h = _SHAPES[0] + x, w13, w2 = _make_inputs(m, d, h) + # Warm up kernel JIT outside the profiled region. + _run(mxfp8_swiglu_mlp_w13, x, w13, w2, fuse_activation=fuse_activation) + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run(mxfp8_swiglu_mlp_w13, x, w13, w2, fuse_activation=fuse_activation) + counts = {} + for evt in prof.key_averages(): + if evt.key in _SWIGLU_OPS or evt.key in _STANDALONE_CAST_OPS: + counts[evt.key] = evt.count + return counts + + +def test_no_standalone_activation_casts_in_fused_mode(): + fused = _op_counts(fuse_activation=True) + unfused = _op_counts(fuse_activation=False) + + # Fused mode: one unified kernel per direction, and the only standalone + # casts are the GEMM-operand casts (fwd: x, w13, w2 rowwise; bwd: go + # rowwise plus w2, w13, x, go colwise). + assert fused.get("torchao::gated_act_mxfp8_forward", 0) == 1 + assert fused.get("torchao::gated_act_mxfp8_backward", 0) == 1 + assert fused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 4 + assert fused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 4 + + # Unfused mode: no unified kernel; the SwiGLU boundary adds exactly one + # rowwise + one colwise standalone cast per direction (h and [dGate|dUp]). + assert unfused.get("torchao::gated_act_mxfp8_forward", 0) == 0 + assert unfused.get("torchao::gated_act_mxfp8_backward", 0) == 0 + assert unfused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 6 + assert unfused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 6 + + +def test_wgrad_with_hp_keeps_fused_activation_casts(): + m, d, h = _SHAPES[0] + x, w13, w2 = _make_inputs(m, d, h) + got = _run( + mxfp8_swiglu_mlp_w13, + x, + w13, + w2, + fuse_activation=True, + wgrad_with_hp=True, + ) + ref = _run( + mxfp8_swiglu_mlp_w13, + x, + w13, + w2, + fuse_activation=False, + wgrad_with_hp=True, + ) + torch.testing.assert_close(got[0], ref[0], rtol=0, atol=0) + for g, r, name in zip(got[1:], ref[1:], ["dx", "dw13", "dw2"]): + sqnr = compute_error(r.float(), g.float()) + assert sqnr >= 50.0, f"{name} SQNR between modes {sqnr} < 50" + + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run( + mxfp8_swiglu_mlp_w13, + x, + w13, + w2, + fuse_activation=True, + wgrad_with_hp=True, + ) + keys = {evt.key for evt in prof.key_averages()} + assert _SWIGLU_OPS <= keys + + +@pytest.mark.parametrize( + "shape", + [ + (100, 256, 768), # M not a multiple of 32: BF16 fallback + (96, 256, 768), # M multiple of 32 but not 128: mx_mm fallback + (256, 256, 704), # H not a multiple of 128: mx_mm fallback + ], +) +def test_unsupported_shapes_fall_back(shape): + m, d, h = shape + x, w13, w2 = _make_inputs(m, d, h) + out, dx, dw13, dw2 = _run(mxfp8_swiglu_mlp_w13, x, w13, w2) + for t, name in [(out, "out"), (dx, "dx"), (dw13, "dw13"), (dw2, "dw2")]: + assert torch.isfinite(t).all(), f"{name} contains non-finite values" + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run(mxfp8_swiglu_mlp_w13, x, w13, w2) + keys = {evt.key for evt in prof.key_averages()} + assert not (keys & _SWIGLU_OPS) + + +@pytest.mark.parametrize("fuse_activation", [True, False]) +def test_no_nans_with_large_inputs(fuse_activation): + m, d, h = _SHAPES[0] + x, w13, w2 = _make_inputs(m, d, h) + with torch.no_grad(): + x.mul_(100.0) + out, dx, dw13, dw2 = _run( + mxfp8_swiglu_mlp_w13, x, w13, w2, fuse_activation=fuse_activation + ) + for t, name in [(out, "out"), (dx, "dx"), (dw13, "dw13"), (dw2, "dw2")]: + assert torch.isfinite(t).all(), f"{name} contains non-finite values" + + +# --------------------------------------------------------------------------- +# Grouped (MoE) composite +# --------------------------------------------------------------------------- + +# Unequal per-expert token groups, all 128-row aligned as the token dispatcher +# guarantees (pad_multiple=128). +_GROUP_SIZES = [256, 128, 384, 256] +_GROUPED_E, _GROUPED_F, _GROUPED_D = 4, 256, 256 + + +def _make_grouped_inputs(sizes=_GROUP_SIZES, f=_GROUPED_F, d=_GROUPED_D, seed=0): + torch.manual_seed(seed) + e = len(sizes) + m = sum(sizes) + offs = torch.tensor(sizes, dtype=torch.int32, device="cuda").cumsum(0).int() + x = torch.randn(m, d, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w13 = torch.randn(e, f, 2, d, dtype=torch.bfloat16, device="cuda") * 0.02 + w2_edf = torch.randn(e, d, f, dtype=torch.bfloat16, device="cuda") * 0.02 + w13.requires_grad_(True) + w2_edf.requires_grad_(True) + return x, w13, w2_edf, offs + + +def _run_grouped(fn, x, w13, w2_edf, offs, **kwargs): + x = x.clone().detach().requires_grad_(True) + w13 = w13.clone().detach().requires_grad_(True) + w2_edf = w2_edf.clone().detach().requires_grad_(True) + out = fn(x, w13, w2_edf.transpose(-2, -1), offs, **kwargs) + torch.manual_seed(1234) + out.backward(torch.randn_like(out)) + return out, x.grad, w13.grad, w2_edf.grad + + +def _bf16_grouped_reference(x, w13, w2_t, offs): + e, f, _, d = w13.shape + w13_packed = w13.transpose(1, 2).reshape(e, 2 * f, d) + gated = torch._grouped_mm(x, w13_packed.transpose(-2, -1), offs=offs) + h = (F.silu(gated[:, :f].float()) * gated[:, f:].float()).to(gated.dtype) + return torch._grouped_mm(h, w2_t, offs=offs) + + +@pytest.mark.parametrize("wgrad_with_hp", [False, True]) +def test_grouped_fused_matches_unfused_mode(wgrad_with_hp): + x, w13, w2_edf, offs = _make_grouped_inputs() + out_f, dx_f, dw13_f, dw2_f = _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=True, + wgrad_with_hp=wgrad_with_hp, + ) + out_u, dx_u, dw13_u, dw2_u = _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=False, + wgrad_with_hp=wgrad_with_hp, + ) + # Forward h quantization is bitwise identical between the modes. + torch.testing.assert_close(out_f, out_u, rtol=0, atol=0) + for got, ref, name in [ + (dx_f, dx_u, "dx"), + (dw13_f, dw13_u, "dw13"), + (dw2_f, dw2_u, "dw2"), + ]: + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= 50.0, f"{name} SQNR between modes {sqnr} < 50" + + +@pytest.mark.parametrize("fuse_activation", [True, False]) +def test_grouped_tracks_references(fuse_activation): + x, w13, w2_edf, offs = _make_grouped_inputs() + out, dx, dw13, dw2 = _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=fuse_activation, + ) + ref = _run_grouped( + lambda *a, **k: _unfused_grouped_mlp(*a, wgrad_with_hp=False), + x, + w13, + w2_edf, + offs, + ) + bf16 = _run_grouped(_bf16_grouped_reference, x, w13, w2_edf, offs) + for got, r, b, name in [ + (out, ref[0], bf16[0], "out"), + (dx, ref[1], bf16[1], "dx"), + (dw13, ref[2], bf16[2], "dw13"), + (dw2, ref[3], bf16[3], "dw2"), + ]: + assert torch.isfinite(got).all(), f"{name} contains non-finite values" + sqnr_q = compute_error(r.float(), got.float()) + assert sqnr_q >= 22.0, f"{name} SQNR vs grouped mx path {sqnr_q} < 22" + sqnr_b = compute_error(b.float(), got.float()) + assert sqnr_b >= 18.0, f"{name} SQNR vs bf16 {sqnr_b} < 18" + + +def test_grouped_padded_rows_do_not_affect_results(): + """Appending a 128-row zero pad block to every expert group must not change + real-row outputs or any gradient (pad rows carry zero upstream grad, as the + un-permute backward guarantees).""" + x, w13, w2_edf, offs = _make_grouped_inputs() + d = x.shape[1] + + xs = list(x.split(_GROUP_SIZES)) + pad = x.new_zeros(128, d) + x_padded = torch.cat([t for g in xs for t in (g, pad)]) + sizes_padded = [s + 128 for s in _GROUP_SIZES] + offs_padded = ( + torch.tensor(sizes_padded, dtype=torch.int32, device="cuda").cumsum(0).int() + ) + + def run(xin, offs_in, sizes): + xin = xin.clone().detach().requires_grad_(True) + w13_ = w13.clone().detach().requires_grad_(True) + w2_ = w2_edf.clone().detach().requires_grad_(True) + out = mxfp8_swiglu_grouped_mlp_w13( + xin, w13_, w2_.transpose(-2, -1), offs_in, fuse_activation=True + ) + torch.manual_seed(1234) + grads = torch.randn(sum(_GROUP_SIZES), out.shape[1], device="cuda").bfloat16() + gsplit = list(grads.split(_GROUP_SIZES)) + if sizes != _GROUP_SIZES: + grads = torch.cat([t for g in gsplit for t in (g, pad)]) + out.backward(grads) + real = torch.cat([t[:s] for t, s in zip(out.split(sizes), _GROUP_SIZES)]) + real_dx = torch.cat( + [t[:s] for t, s in zip(xin.grad.split(sizes), _GROUP_SIZES)] + ) + return real, real_dx, w13_.grad, w2_.grad + + out_a, dx_a, dw13_a, dw2_a = run(x, offs, _GROUP_SIZES) + out_b, dx_b, dw13_b, dw2_b = run(x_padded, offs_padded, sizes_padded) + torch.testing.assert_close(out_a, out_b, rtol=0, atol=0) + torch.testing.assert_close(dx_a, dx_b, rtol=0, atol=0) + torch.testing.assert_close(dw13_a, dw13_b, rtol=0, atol=0) + torch.testing.assert_close(dw2_a, dw2_b, rtol=0, atol=0) + + +def test_grouped_global_tail_padding_is_inert(): + """The dispatcher may pad the token buffer globally past offs[-1]; those + tail rows belong to no expert group and must not affect real-row outputs + or any gradient (and must not crash the scale-tile gather).""" + x, w13, w2_edf, offs = _make_grouped_inputs() + m, d = x.shape + + def run(xin, total_rows): + xin = xin.clone().detach().requires_grad_(True) + w13_ = w13.clone().detach().requires_grad_(True) + w2_ = w2_edf.clone().detach().requires_grad_(True) + out = mxfp8_swiglu_grouped_mlp_w13( + xin, w13_, w2_.transpose(-2, -1), offs, fuse_activation=True + ) + torch.manual_seed(1234) + grads = torch.randn(m, out.shape[1], device="cuda").bfloat16() + if total_rows > m: + grads = torch.cat([grads, grads.new_zeros(total_rows - m, out.shape[1])]) + out.backward(grads) + return out[:m], xin.grad[:m], w13_.grad, w2_.grad + + ref = run(x, m) + tail = run(torch.cat([x.detach(), x.new_zeros(128, d)]), m + 128) + for a, b, name in zip(ref, tail, ["out", "dx", "dw13", "dw2"]): + torch.testing.assert_close(a, b, rtol=0, atol=0) + + +@pytest.mark.parametrize("unbacked_m", [False, True]) +@pytest.mark.parametrize("wgrad_with_hp", [False, True]) +def test_grouped_compile(unbacked_m, wgrad_with_hp): + x, w13, w2_edf, offs = _make_grouped_inputs() + + def fn(x, w13, w2_edf): + return mxfp8_swiglu_grouped_mlp_w13( + x, + w13, + w2_edf.transpose(-2, -1), + offs, + fuse_activation=True, + wgrad_with_hp=wgrad_with_hp, + ) + + eager = _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=True, + wgrad_with_hp=wgrad_with_hp, + ) + compiled_fn = torch.compile(fn, fullgraph=True) + + xc = x.clone().detach().requires_grad_(True) + if unbacked_m: + # Routing-dependent M as EP token dispatch produces it: unbacked under + # compile, so the support checks must defer to runtime asserts instead + # of failing at trace time. + torch._dynamo.decorators.mark_unbacked(xc, 0) + w13c = w13.clone().detach().requires_grad_(True) + w2c = w2_edf.clone().detach().requires_grad_(True) + out = compiled_fn(xc, w13c, w2c) + torch.manual_seed(1234) + out.backward(torch.randn_like(out)) + + torch.testing.assert_close(out, eager[0], rtol=0, atol=0) + for got, ref, name in [ + (xc.grad, eager[1], "dx"), + (w13c.grad, eager[2], "dw13"), + (w2c.grad, eager[3], "dw2"), + ]: + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= 50.0, f"compiled {name} SQNR {sqnr} < 50" + + +def _grouped_op_counts(fuse_activation): + x, w13, w2_edf, offs = _make_grouped_inputs() + _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=fuse_activation, + ) + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=fuse_activation, + ) + counts = {} + for evt in prof.key_averages(): + if evt.key in _SWIGLU_OPS or evt.key in _STANDALONE_CAST_OPS: + counts[evt.key] = evt.count + return counts + + +def test_grouped_no_standalone_activation_casts_in_fused_mode(): + fused = _grouped_op_counts(fuse_activation=True) + unfused = _grouped_op_counts(fuse_activation=False) + + # Fused mode: the only 2D cutedsl casts are the GEMM-operand rowwise casts + # of x (fwd) and grad_out (bwd); wgrad colwise casts use the CUDA dim1 + # kernel, mirroring the existing grouped wgrad path. + assert fused.get("torchao::gated_act_mxfp8_forward", 0) == 1 + assert fused.get("torchao::gated_act_mxfp8_backward", 0) == 1 + assert fused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 2 + assert fused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 0 + + # Unfused mode: the SwiGLU boundary adds one rowwise + one colwise + # standalone cast per direction (h and [dGate | dUp]). + assert unfused.get("torchao::gated_act_mxfp8_forward", 0) == 0 + assert unfused.get("torchao::gated_act_mxfp8_backward", 0) == 0 + assert unfused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 4 + assert unfused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 2 + + +def test_grouped_hp_wgrad_keeps_fused_activation_casts(): + x, w13, w2_edf, offs = _make_grouped_inputs() + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run_grouped( + mxfp8_swiglu_grouped_mlp_w13, + x, + w13, + w2_edf, + offs, + fuse_activation=True, + wgrad_with_hp=True, + ) + keys = {evt.key for evt in prof.key_averages()} + assert _SWIGLU_OPS <= keys + + +def test_grouped_unsupported_shape_falls_back(): + # F a multiple of 32 but not 128: falls back, still finite. + x, w13, w2_edf, offs = _make_grouped_inputs(f=192) + res = _run_grouped(mxfp8_swiglu_grouped_mlp_w13, x, w13, w2_edf, offs) + for t, name in zip(res, ["out", "dx", "dw13", "dw2"]): + assert torch.isfinite(t).all(), f"{name} contains non-finite values" + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run_grouped(mxfp8_swiglu_grouped_mlp_w13, x, w13, w2_edf, offs) + keys = {evt.key for evt in prof.key_averages()} + assert not (keys & _SWIGLU_OPS) diff --git a/torchtitan/overrides/fused_swiglu.py b/torchtitan/overrides/fused_swiglu.py index 20cd2d3c2e..2d4980aa2f 100644 --- a/torchtitan/overrides/fused_swiglu.py +++ b/torchtitan/overrides/fused_swiglu.py @@ -456,7 +456,7 @@ def __init__(self, config: Config): def forward(self, x: torch.Tensor) -> torch.Tensor: if self.mxfp8_fused and not isinstance(x, DTensor): - from torchao.prototype.moe_training import mxfp8_swiglu_mlp_w13 + from torchtitan.overrides.swiglu_mlp import mxfp8_swiglu_mlp_w13 output = mxfp8_swiglu_mlp_w13( x, @@ -588,7 +588,7 @@ def forward( offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) if self.mxfp8_fused: - from torchao.prototype.moe_training import mxfp8_swiglu_grouped_mlp_w13 + from torchtitan.overrides.swiglu_mlp import mxfp8_swiglu_grouped_mlp_w13 return mxfp8_swiglu_grouped_mlp_w13( x_RD.bfloat16(), diff --git a/torchtitan/overrides/swiglu_mlp.py b/torchtitan/overrides/swiglu_mlp.py new file mode 100644 index 0000000000..815cd6b046 --- /dev/null +++ b/torchtitan/overrides/swiglu_mlp.py @@ -0,0 +1,544 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Composite MXFP8 SwiGLU MLP for a fused w13 projection. + +One autograd function covers the full dense MLP + + x -> MXFP8 w13 GEMM -> [gate | up] -> silu(gate) * up -> MXFP8 w2 GEMM + +with both directions of every quantization done by the CuTeDSL kernels. The +``fuse_activation`` flag selects how the activation boundary is quantized: + +* ``True``: the unified SwiGLU+MXFP8 kernel produces the rowwise and colwise + MXFP8 copies directly; the BF16 activation ``h`` is never written to global + memory. +* ``False``: ``h`` (forward) and ``[dGate | dUp]`` (backward) are materialized + in BF16 and quantized by the standalone 1x32 / 32x1 CuTeDSL kernels. + +Everything outside that boundary -- the w13/w2 GEMMs, their input, weight and +gradient casts -- is byte-for-byte identical between the two modes, so an A/B +comparison isolates the activation+quantization implementation. + +The composite supports both quantized and high-precision weight gradients. +With ``wgrad_with_hp=True`` the fused kernels still provide the forward and +dgrad activation casts, while the BF16 activation values needed by the two +weight-gradient GEMMs are recomputed from the saved gated projection. Shapes +the CuTeDSL kernels cannot handle still fall back to the existing per-GEMM +``mx_mm`` path (or plain BF16 as a last resort) instead of asserting. +""" + +import torch +import torch.nn.functional as F +from torch.distributed.tensor import DTensor + +from torchao.prototype.moe_training.kernels.mxfp8 import ( + triton_mx_block_rearrange_2d_K_groups, +) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( + _mxfp8_cutedsl_kernels_available, + mxfp8_quantize_2d_1x32_cutedsl, + mxfp8_quantize_2d_32x1_cutedsl, + gated_act_mxfp8_backward, + gated_act_mxfp8_forward, +) +from torchao.prototype.moe_training.mxfp8_grouped_mm import ( + _compute_dgrad_sm100, + _compute_fwd_sm100, + _to_mxfp8_then_scaled_grouped_mm, +) +from torchao.prototype.moe_training.mxfp8_linear import _to_mxfp8_then_scaled_mm +from torchao.prototype.mx_formats.config import ( + MXFP8Dim1CastKernelChoice, + ScaleCalculationMode, +) +from torchao.prototype.mx_formats.mx_tensor import MXTensor +from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper +from torchao.quantization.quantize_.common.kernel_preference import KernelPreference + +__all__ = ["mxfp8_swiglu_mlp_w13", "mxfp8_swiglu_grouped_mlp_w13"] + +_BLOCK_SIZE = 32 +_ELEM_DTYPE = torch.float8_e4m3fn +_KERNEL_PREFERENCE = KernelPreference.AUTO +_SCALE_MODE = ScaleCalculationMode.RCEIL +_INT32_MAX = 2**31 - 1 + + +def _wrap_rowwise(qdata, scales, orig_dtype): + return MXTensor.from_qdata_and_scales( + qdata, + scales, + orig_dtype, + block_size=_BLOCK_SIZE, + kernel_preference=_KERNEL_PREFERENCE, + is_swizzled_scales=True, + ) + + +def _wrap_colwise(qdata, scales, orig_dtype): + # Colwise kernel outputs are (M, N) with strides (1, M); wrapping the + # transpose keeps qdata row-major, which torch.mm's MXFP8 dispatch + # requires. The flat 1D blocked scales are unaffected by the transpose. + return _wrap_rowwise(qdata.t(), scales, orig_dtype) + + +def _mx_rowwise(t): + qdata, scales = mxfp8_quantize_2d_1x32_cutedsl(t, scaling_mode=_SCALE_MODE.value) + return _wrap_rowwise(qdata, scales, t.dtype) + + +def _mx_colwise(t): + # Returns the MXTensor for t.t() quantized along t's rows (32x1 blocks). + return _to_mxfp8_dim1_kernel_wrapper( + t, + _BLOCK_SIZE, + _ELEM_DTYPE, + t.dtype, + _KERNEL_PREFERENCE, + MXFP8Dim1CastKernelChoice.CUTEDSL, + _SCALE_MODE, + ) + + +def _pack_w13(w13): + # (H, 2, D) with w13[:, 0] = gate and w13[:, 1] = up, packed to (2H, D) + # with all gate rows first -- the layout whose GEMM output feeds the + # SwiGLU kernel's [gate | up] contract. + hidden, _, dim = w13.shape + return w13.transpose(0, 1).reshape(2 * hidden, dim).contiguous() + + +def _swiglu_forward_hp(gated): + k = gated.shape[1] // 2 + return (F.silu(gated[:, :k].float()) * gated[:, k:].float()).to(gated.dtype) + + +def _empty_mxfp8_outputs(t): + return ( + t.new_empty(0, dtype=torch.float8_e4m3fn), + t.new_empty(0, dtype=torch.float8_e8m0fnu), + ) + + +def _swiglu_forward_casts(gated, fuse_activation, colwise): + if fuse_activation: + return gated_act_mxfp8_forward(gated, rowwise=True, colwise=colwise) + h = _swiglu_forward_hp(gated) + h_rw, hs_rw = mxfp8_quantize_2d_1x32_cutedsl(h, scaling_mode=_SCALE_MODE.value) + if colwise: + h_cw, hs_cw = mxfp8_quantize_2d_32x1_cutedsl(h, scaling_mode=_SCALE_MODE.value) + else: + h_cw, hs_cw = _empty_mxfp8_outputs(gated) + return h_rw, h_cw, hs_rw, hs_cw + + +def _swiglu_backward_hp(grad_h, gated): + k = gated.shape[1] // 2 + gate = gated[:, :k].float() + up = gated[:, k:].float() + grad_h_f = grad_h.float() + # Same evaluation order as the unified kernel (which contracts `deriv` + # into one FMA), so the two modes differ only in sigmoid lowering and + # that contraction, not in association. + sigmoid_gate = torch.sigmoid(gate) + silu = gate * sigmoid_gate + deriv = gate * (1.0 - sigmoid_gate) + 1.0 + return torch.cat( + [ + ((grad_h_f * up) * (sigmoid_gate * deriv)).to(gated.dtype), + (grad_h_f * silu).to(gated.dtype), + ], + dim=1, + ) + + +def _swiglu_backward_casts(grad_h, gated, fuse_activation, colwise): + if fuse_activation: + return gated_act_mxfp8_backward(grad_h, gated, rowwise=True, colwise=colwise) + d = _swiglu_backward_hp(grad_h, gated) + d_rw, ds_rw = mxfp8_quantize_2d_1x32_cutedsl(d, scaling_mode=_SCALE_MODE.value) + if colwise: + d_cw, ds_cw = mxfp8_quantize_2d_32x1_cutedsl(d, scaling_mode=_SCALE_MODE.value) + else: + d_cw, ds_cw = _empty_mxfp8_outputs(gated) + return d_rw, d_cw, ds_rw, ds_cw + + +@torch._dynamo.allow_in_graph +class _MXFP8SwiGLUMLP(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w13, w2, fuse_activation, wgrad_with_hp): + x2d = x.reshape(-1, x.shape[-1]).contiguous() + w13_packed = _pack_w13(w13) + gated = torch.mm(_mx_rowwise(x2d), _mx_rowwise(w13_packed).t()) + h_rw, h_cw, hs_rw, hs_cw = _swiglu_forward_casts( + gated, fuse_activation, colwise=not wgrad_with_hp + ) + out = torch.mm(_wrap_rowwise(h_rw, hs_rw, x2d.dtype), _mx_rowwise(w2).t()) + ctx.save_for_backward(x2d, w13_packed, w2, gated, h_cw, hs_cw) + ctx.fuse_activation = fuse_activation + ctx.wgrad_with_hp = wgrad_with_hp + ctx.x_shape = x.shape + return out.reshape(*x.shape[:-1], out.shape[-1]) + + @staticmethod + def backward(ctx, grad_out): + x2d, w13_packed, w2, gated, h_cw, hs_cw = ctx.saved_tensors + hidden = w13_packed.shape[0] // 2 + go = grad_out.reshape(-1, grad_out.shape[-1]).contiguous() + grad_h = torch.mm(_mx_rowwise(go), _mx_colwise(w2).t()) + d_rw, d_cw, ds_rw, ds_cw = _swiglu_backward_casts( + grad_h, + gated, + ctx.fuse_activation, + colwise=not ctx.wgrad_with_hp, + ) + grad_x = torch.mm( + _wrap_rowwise(d_rw, ds_rw, go.dtype), _mx_colwise(w13_packed).t() + ) + if ctx.wgrad_with_hp: + # The fused casts feed dgrad above. Recompute the BF16 activation + # boundary only for the two HP wgrad GEMMs, avoiding a forward HBM + # write of h while preserving the recipe's BF16 wgrad semantics. + d_hp = _swiglu_backward_hp(grad_h, gated) + h_hp = _swiglu_forward_hp(gated) + grad_w13_packed = torch.mm(d_hp.t(), x2d) + grad_w2 = torch.mm(go.t(), h_hp) + else: + grad_w13_packed = torch.mm( + _wrap_colwise(d_cw, ds_cw, go.dtype), _mx_colwise(x2d).t() + ) + grad_w2 = torch.mm( + _mx_colwise(go), _wrap_colwise(h_cw, hs_cw, go.dtype).t() + ) + grad_w13 = grad_w13_packed.view(2, hidden, -1).transpose(0, 1).contiguous() + return grad_x.reshape(ctx.x_shape), grad_w13, grad_w2, None, None + + +def _fused_path_ok(x, w13, w2): + if not _mxfp8_cutedsl_kernels_available: + return False + if isinstance(x, DTensor) or isinstance(w13, DTensor) or isinstance(w2, DTensor): + return False + if not x.is_cuda: + return False + if ( + x.dtype != torch.bfloat16 + or w13.dtype != torch.bfloat16 + or w2.dtype != torch.bfloat16 + ): + return False + if w13.ndim != 3 or w13.shape[1] != 2 or w2.ndim != 2: + return False + hidden, _, dim = w13.shape + n = w2.shape[0] + m = x.numel() // x.shape[-1] + if x.shape[-1] != dim or w2.shape[1] != hidden: + return False + if m % 128 != 0 or hidden % 128 != 0 or dim % 128 != 0 or n % 128 != 0: + return False + # 32-bit index-math limit over BOTH A/B arms: the unified kernel's input + # layout reaches element 2*hidden*m - hidden - 1, but the unfused arm's + # standalone casts of the (m, 2*hidden) backward tensor reach + # 2*hidden*m - 1, and those kernels do not validate. Gate on the max so + # the two arms accept identical shapes. + if 2 * hidden * m - 1 > _INT32_MAX: + return False + return True + + +def _mx_mm_path_ok(x, w13, w2): + if isinstance(x, DTensor) or isinstance(w13, DTensor) or isinstance(w2, DTensor): + return False + if not x.is_cuda or x.dtype != torch.bfloat16: + return False + hidden, _, dim = w13.shape + m = x.numel() // x.shape[-1] + return m % 32 == 0 and hidden % 32 == 0 and dim % 32 == 0 and w2.shape[0] % 32 == 0 + + +def _unfused_mlp(x, w13, w2, wgrad_with_hp): + if _mx_mm_path_ok(x, w13, w2): + gate_up = _to_mxfp8_then_scaled_mm( + x, _pack_w13(w13), _KERNEL_PREFERENCE, _SCALE_MODE, wgrad_with_hp + ) + gate, up = gate_up.chunk(2, dim=-1) + h = F.silu(gate) * up + return _to_mxfp8_then_scaled_mm( + h, w2, _KERNEL_PREFERENCE, _SCALE_MODE, wgrad_with_hp + ) + gate, up = torch.einsum("...d,hgd->...hg", x, w13).unbind(-1) + return F.linear(F.silu(gate.float()).to(x.dtype) * up, w2) + + +def mxfp8_swiglu_mlp_w13( + x, w13, down_weight, *, fuse_activation=True, wgrad_with_hp=False +): + """Dense MXFP8 SwiGLU MLP with a fused (H, 2, D) w13 weight. + + Args: + x: BF16 input of shape (..., D). + w13: BF16 fused gate/up weight of shape (H, 2, D); w13[:, 0] is the + gate (w1) and w13[:, 1] the up (w3) projection. + down_weight: BF16 down-projection weight of shape (D_out, H). + fuse_activation: quantize the SwiGLU boundary with the unified + SwiGLU+MXFP8 kernel instead of standalone BF16 + cast kernels. + wgrad_with_hp: compute the two weight gradients with BF16 GEMMs while + retaining the fused forward/dgrad activation casts. + + Returns: + BF16 tensor of shape (..., D_out). + """ + if not _fused_path_ok(x, w13, down_weight): + return _unfused_mlp(x, w13, down_weight, wgrad_with_hp) + return _MXFP8SwiGLUMLP.apply(x, w13, down_weight, fuse_activation, wgrad_with_hp) + + +def _pack_w13_grouped(w13): + # (E, F, 2, D) with w13[:, :, 0] = gate and w13[:, :, 1] = up, packed to + # (E, 2F, D) with all gate rows first per expert (the [gate | up] layout + # the SwiGLU kernel consumes). The reshape of the transposed view copies. + e, f, _, d = w13.shape + return w13.transpose(1, 2).reshape(e, 2 * f, d) + + +def _reblock_scales_k_groups(scales, n_rows, m_total, offs): + # The CuTeDSL kernels emit blocked scales in full-tensor row-block-major + # tile order; a 2d-2d grouped GEMM contracting over tokens needs the tiles + # regrouped per token group (row-block-major within each group). With every + # group a multiple of 128 rows the two layouts hold identical (128, 4) + # tiles, so this is a pure tile gather. + rb = n_rows // 128 + cb = m_total // 128 + ends = (offs // 128).long() + starts = torch.cat([ends.new_zeros(1), ends[:-1]]) + sizes = (ends - starts).clamp(min=1) + t = torch.arange(rb * cb, device=scales.device) + # The dispatcher may pad the token buffer globally past offs[-1]; tiles in + # that tail belong to no group and are never read by the grouped GEMM, so + # clamping them anywhere in bounds is enough to keep the gather valid. + g = torch.searchsorted(ends * rb, t, right=True).clamp(max=ends.numel() - 1) + local = t - starts[g] * rb + src = ((local // sizes[g]) * cb + starts[g] + local % sizes[g]).clamp( + max=rb * cb - 1 + ) + return scales.view(rb * cb, 512)[src].view(n_rows, -1) + + +def _wgrad_k_groups(a_qdata, a_scales, b, offs, out_dtype): + # grad[e] = a[start:end].T @ b[start:end] for each token group. `a` arrives + # colwise-quantized from the SwiGLU boundary ((M, Ka) with strides (1, M) + # plus flat blocked scales); `b` gets the same GEMM-operand colwise cast the + # existing grouped wgrad path uses. + m, ka = a_qdata.shape + b_t_mx = _to_mxfp8_dim1_kernel_wrapper( + b, + _BLOCK_SIZE, + _ELEM_DTYPE, + b.dtype, + _KERNEL_PREFERENCE, + MXFP8Dim1CastKernelChoice.CUDA, + _SCALE_MODE, + ) + b_scales = triton_mx_block_rearrange_2d_K_groups(b_t_mx.scale, offs // _BLOCK_SIZE) + a_scales_2d = _reblock_scales_k_groups(a_scales, ka, m, offs) + return torch._scaled_grouped_mm( + a_qdata.t(), + b_t_mx.qdata.transpose(-2, -1), + a_scales_2d, + b_scales, + offs=offs, + out_dtype=out_dtype, + ) + + +@torch._dynamo.allow_in_graph +class _MXFP8SwiGLUGroupedMLP(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w13, w2_t, offs, fuse_activation, wgrad_with_hp): + x = x.contiguous() + w13_packed = _pack_w13_grouped(w13) + gated = _compute_fwd_sm100( + x, w13_packed.transpose(-2, -1), offs, _BLOCK_SIZE, x.dtype, _SCALE_MODE + ) + h_rw, h_cw, hs_rw, hs_cw = _swiglu_forward_casts( + gated, fuse_activation, colwise=not wgrad_with_hp + ) + out = _compute_fwd_sm100( + _wrap_rowwise(h_rw, hs_rw, x.dtype), + w2_t, + offs, + _BLOCK_SIZE, + x.dtype, + _SCALE_MODE, + ) + ctx.save_for_backward(x, w13_packed, w2_t, offs, gated, h_cw, hs_cw) + ctx.fuse_activation = fuse_activation + ctx.wgrad_with_hp = wgrad_with_hp + return out + + @staticmethod + def backward(ctx, grad_out): + x, w13_packed, w2_t, offs, gated, h_cw, hs_cw = ctx.saved_tensors + e, two_f, d = w13_packed.shape + go = grad_out.contiguous() + grad_h = _compute_dgrad_sm100( + go, w2_t, offs, _BLOCK_SIZE, go.dtype, _SCALE_MODE + ) + d_rw, d_cw, ds_rw, ds_cw = _swiglu_backward_casts( + grad_h, + gated, + ctx.fuse_activation, + colwise=not ctx.wgrad_with_hp, + ) + grad_x = _compute_dgrad_sm100( + _wrap_rowwise(d_rw, ds_rw, go.dtype), + w13_packed.transpose(-2, -1), + offs, + _BLOCK_SIZE, + go.dtype, + _SCALE_MODE, + ) + if ctx.wgrad_with_hp: + d_hp = _swiglu_backward_hp(grad_h, gated) + h_hp = _swiglu_forward_hp(gated) + grad_w13_packed = torch._grouped_mm( + d_hp.t(), x, offs=offs, out_dtype=go.dtype + ) + grad_w2_t = torch._grouped_mm(h_hp.t(), go, offs=offs, out_dtype=go.dtype) + else: + grad_w13_packed = _wgrad_k_groups(d_cw, ds_cw, x, offs, go.dtype) + grad_w2_t = _wgrad_k_groups(h_cw, hs_cw, go, offs, go.dtype) + grad_w13 = grad_w13_packed.view(e, 2, two_f // 2, d).transpose(1, 2) + return grad_x, grad_w13, grad_w2_t, None, None, None + + +def _grouped_path_ok(x, w13, w2_t, offs): + if not _mxfp8_cutedsl_kernels_available: + return False + if any(isinstance(t, DTensor) for t in (x, w13, w2_t)): + return False + if not x.is_cuda or x.ndim != 2: + return False + if ( + x.dtype != torch.bfloat16 + or w13.dtype != torch.bfloat16 + or w2_t.dtype != torch.bfloat16 + ): + return False + if w13.ndim != 4 or w13.shape[2] != 2 or w2_t.ndim != 3: + return False + e, f, _, d = w13.shape + m = x.shape[0] + d_out = w2_t.shape[2] + if x.shape[1] != d or w2_t.shape[:2] != (e, f) or offs.shape != (e,): + return False + if f % 128 != 0 or d % 128 != 0 or d_out % 128 != 0: + return False + # Group boundaries must additionally be 128-row aligned (the token + # dispatcher's pad_multiple guarantees it); checking offs here would sync. + # M is routing-dependent under compile (an unbacked SymInt, which type + # tests cannot tell apart from int inside traced code), so the M + # conditions use identity tests: literal bools keep the trace-time + # fallback, symbolic ones become deferred runtime asserts. The m >= 128 + # and m % 32 forms are redundant with m % 128 (plus non-emptiness) but + # must be recorded separately: downstream cast-kernel wrappers and GEMM + # metas check exactly those forms, and the symbolic engine resolves them + # by expression match / value range, not by deriving them from mod-128. + # The last condition is the 32-bit index-math limit over BOTH A/B arms + # (the unfused arm's standalone casts of the (m, 2f) backward tensor + # reach element 2*f*m - 1, slightly past the unified kernel's own input + # bound, and those kernels do not validate). + for cond in ( + m >= 128, + m % 128 == 0, + m % 32 == 0, + 2 * f * m - 1 <= _INT32_MAX, + ): + if cond is False: + return False + if cond is not True: + torch._check(cond) + return True + + +def _unfused_grouped_mlp(x, w13, w2_t, offs, wgrad_with_hp): + # The per-GEMM grouped mx path's cast kernels need every dim to be a + # multiple of 128, so it only serves the wgrad_with_hp (and non-shape + # guard) fallbacks; unsupported shapes drop to BF16 grouped GEMMs. + e, f, _, d = w13.shape + m = x.shape[0] + mx_ok = ( + x.is_cuda + and x.dtype == torch.bfloat16 + and f % 128 == 0 + and d % 128 == 0 + and w2_t.shape[2] % 128 == 0 + ) + if mx_ok: + # Same identity-test treatment of the routing-dependent M as + # _grouped_path_ok, for when this fallback is itself compiled. + m_ok = m % 128 == 0 + if m_ok is False: + mx_ok = False + elif m_ok is not True: + torch._check(m_ok) + if mx_ok: + gated = _to_mxfp8_then_scaled_grouped_mm( + x.contiguous(), + _pack_w13_grouped(w13).transpose(-2, -1), + offs, + kernel_preference=_KERNEL_PREFERENCE, + wgrad_with_hp=wgrad_with_hp, + scale_calculation_mode=_SCALE_MODE, + ) + h = (F.silu(gated[:, :f].float()) * gated[:, f:].float()).to(gated.dtype) + return _to_mxfp8_then_scaled_grouped_mm( + h, + w2_t, + offs, + kernel_preference=_KERNEL_PREFERENCE, + wgrad_with_hp=wgrad_with_hp, + scale_calculation_mode=_SCALE_MODE, + ) + gated = torch._grouped_mm(x, _pack_w13_grouped(w13).transpose(-2, -1), offs=offs) + h = (F.silu(gated[:, :f].float()) * gated[:, f:].float()).to(gated.dtype) + return torch._grouped_mm(h, w2_t, offs=offs) + + +def mxfp8_swiglu_grouped_mlp_w13( + x, w13, down_weight_t, offs, *, fuse_activation=True, wgrad_with_hp=False +): + """Grouped-expert MXFP8 SwiGLU MLP with a fused (E, F, 2, D) w13 weight. + + Args: + x: BF16 token rows of shape (M, D) in expert-major order, with every + expert's group padded to a multiple of 128 rows (padded rows must + be zero) so 32x1 scale blocks never cross group boundaries. + w13: BF16 fused gate/up weight of shape (E, F, 2, D); w13[:, :, 0] is + the gate and w13[:, :, 1] the up projection. + down_weight_t: BF16 down-projection weight of shape (E, F, D_out) in + per-expert column-major layout (a transposed view of (E, D_out, F)). + offs: int32 group end offsets of shape (E,), each a multiple of 128. + fuse_activation: quantize the SwiGLU boundary with the unified + SwiGLU+MXFP8 kernel instead of standalone BF16 + cast kernels. + wgrad_with_hp: compute the two grouped weight gradients with BF16 + GEMMs while retaining the fused forward/dgrad activation casts. + + Returns: + BF16 tensor of shape (M, D_out). + """ + if not _grouped_path_ok(x, w13, down_weight_t, offs): + return _unfused_grouped_mlp(x, w13, down_weight_t, offs, wgrad_with_hp) + return _MXFP8SwiGLUGroupedMLP.apply( + x, + w13, + down_weight_t, + offs, + fuse_activation, + wgrad_with_hp, + ) From 025c794f2ed86c1890b8d89735ab51daa257f070 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 22:11:35 -0700 Subject: [PATCH 4/4] Follow torchao's gated-act op move and rename torchao moved the fused gated-activation custom ops out of quant.py into the kernel module and renamed them with the cutedsl infix (gated_act_mxfp8_cutedsl_{forward,backward}). The composite now imports the wrappers lazily at the two fused-path call sites -- the kernel module imports the CuTe DSL runtime at module scope, and the unfused fallback must keep working without it. Trace-count assertions updated to the new op names. 52/52 swiglu unit tests green against torchao 39db5297. Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_swiglu_mlp.py | 20 ++++++++++---------- torchtitan/overrides/swiglu_mlp.py | 18 ++++++++++++++---- 2 files changed, 24 insertions(+), 14 deletions(-) diff --git a/tests/unit_tests/test_swiglu_mlp.py b/tests/unit_tests/test_swiglu_mlp.py index 3c8e38f65d..a3aacdd1e8 100644 --- a/tests/unit_tests/test_swiglu_mlp.py +++ b/tests/unit_tests/test_swiglu_mlp.py @@ -36,8 +36,8 @@ ] _SWIGLU_OPS = { - "torchao::gated_act_mxfp8_forward", - "torchao::gated_act_mxfp8_backward", + "torchao::gated_act_mxfp8_cutedsl_forward", + "torchao::gated_act_mxfp8_cutedsl_backward", } _STANDALONE_CAST_OPS = { "torchao::mxfp8_quantize_2d_1x32_cutedsl", @@ -197,15 +197,15 @@ def test_no_standalone_activation_casts_in_fused_mode(): # Fused mode: one unified kernel per direction, and the only standalone # casts are the GEMM-operand casts (fwd: x, w13, w2 rowwise; bwd: go # rowwise plus w2, w13, x, go colwise). - assert fused.get("torchao::gated_act_mxfp8_forward", 0) == 1 - assert fused.get("torchao::gated_act_mxfp8_backward", 0) == 1 + assert fused.get("torchao::gated_act_mxfp8_cutedsl_forward", 0) == 1 + assert fused.get("torchao::gated_act_mxfp8_cutedsl_backward", 0) == 1 assert fused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 4 assert fused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 4 # Unfused mode: no unified kernel; the SwiGLU boundary adds exactly one # rowwise + one colwise standalone cast per direction (h and [dGate|dUp]). - assert unfused.get("torchao::gated_act_mxfp8_forward", 0) == 0 - assert unfused.get("torchao::gated_act_mxfp8_backward", 0) == 0 + assert unfused.get("torchao::gated_act_mxfp8_cutedsl_forward", 0) == 0 + assert unfused.get("torchao::gated_act_mxfp8_cutedsl_backward", 0) == 0 assert unfused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 6 assert unfused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 6 @@ -535,15 +535,15 @@ def test_grouped_no_standalone_activation_casts_in_fused_mode(): # Fused mode: the only 2D cutedsl casts are the GEMM-operand rowwise casts # of x (fwd) and grad_out (bwd); wgrad colwise casts use the CUDA dim1 # kernel, mirroring the existing grouped wgrad path. - assert fused.get("torchao::gated_act_mxfp8_forward", 0) == 1 - assert fused.get("torchao::gated_act_mxfp8_backward", 0) == 1 + assert fused.get("torchao::gated_act_mxfp8_cutedsl_forward", 0) == 1 + assert fused.get("torchao::gated_act_mxfp8_cutedsl_backward", 0) == 1 assert fused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 2 assert fused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 0 # Unfused mode: the SwiGLU boundary adds one rowwise + one colwise # standalone cast per direction (h and [dGate | dUp]). - assert unfused.get("torchao::gated_act_mxfp8_forward", 0) == 0 - assert unfused.get("torchao::gated_act_mxfp8_backward", 0) == 0 + assert unfused.get("torchao::gated_act_mxfp8_cutedsl_forward", 0) == 0 + assert unfused.get("torchao::gated_act_mxfp8_cutedsl_backward", 0) == 0 assert unfused.get("torchao::mxfp8_quantize_2d_1x32_cutedsl", 0) == 4 assert unfused.get("torchao::mxfp8_quantize_2d_32x1_cutedsl", 0) == 2 diff --git a/torchtitan/overrides/swiglu_mlp.py b/torchtitan/overrides/swiglu_mlp.py index 815cd6b046..b06eaca317 100644 --- a/torchtitan/overrides/swiglu_mlp.py +++ b/torchtitan/overrides/swiglu_mlp.py @@ -42,8 +42,6 @@ _mxfp8_cutedsl_kernels_available, mxfp8_quantize_2d_1x32_cutedsl, mxfp8_quantize_2d_32x1_cutedsl, - gated_act_mxfp8_backward, - gated_act_mxfp8_forward, ) from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _compute_dgrad_sm100, @@ -126,7 +124,13 @@ def _empty_mxfp8_outputs(t): def _swiglu_forward_casts(gated, fuse_activation, colwise): if fuse_activation: - return gated_act_mxfp8_forward(gated, rowwise=True, colwise=colwise) + # Lazy: the kernel module imports the CuTe DSL runtime at module + # scope; the unfused fallback must work without it. + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_gated_act_mxfp8 import ( + gated_act_mxfp8_cutedsl_forward, + ) + + return gated_act_mxfp8_cutedsl_forward(gated, rowwise=True, colwise=colwise) h = _swiglu_forward_hp(gated) h_rw, hs_rw = mxfp8_quantize_2d_1x32_cutedsl(h, scaling_mode=_SCALE_MODE.value) if colwise: @@ -158,7 +162,13 @@ def _swiglu_backward_hp(grad_h, gated): def _swiglu_backward_casts(grad_h, gated, fuse_activation, colwise): if fuse_activation: - return gated_act_mxfp8_backward(grad_h, gated, rowwise=True, colwise=colwise) + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_gated_act_mxfp8 import ( + gated_act_mxfp8_cutedsl_backward, + ) + + return gated_act_mxfp8_cutedsl_backward( + grad_h, gated, rowwise=True, colwise=colwise + ) d = _swiglu_backward_hp(grad_h, gated) d_rw, ds_rw = mxfp8_quantize_2d_1x32_cutedsl(d, scaling_mode=_SCALE_MODE.value) if colwise: