From 6e2623b4fd0e5bc2c401d22225bc8f69de414cd6 Mon Sep 17 00:00:00 2001 From: Yihua Wei Date: Thu, 5 Nov 2026 00:33:59 +0000 Subject: [PATCH] feat(cudnn_frost): SM120 launch-time SwiGLU parameters, gpt-oss H/I 2880, kernels re-exported from cuDNN Frontend 886d7176 The cuDNN Frost MoE backends rejected every model whose SwiGLU is not the default one: DeepSeek-V4/V4.1-Flash, its DSpark draft head and GLM-5.3-Flash clamp at limit 10, and gpt-oss and MiniMax-M3 use SwiGLU(alpha=1.702, beta=1, limit=7). On SM120 they now run, in all four dtypes. - swiglu_param: SM120 sources whose epilogue reads alpha, beta, limit and -limit as FP32 scalar kernel inputs (FlashInfer's SwiGLU contract: g = min(gate, limit), u = clamp(up, -limit, limit), g * sigmoid(alpha * g) * (u + beta)), so one kernel serves every model. The default SwiGLU() keeps its swiglu sources. The native executors write the scalars into their scalar slots during routing, so CUDA graph replays keep them. SM107 packages no swiglu_param sources, and its runner declines such activations instead of dropping the parameters. - gpt-oss geometry (H = I = 2880): SM120 kernels take N/K in whole 32-element blocks, so the SM120 catalog contract needs multiples of 32. BF16 admits H/I in 32s and NVFP4 in 64s on SM120. MXFP8 and MXFP8 x MXFP4 keep 128s: their canonical weight scales are int32-packed per 128 K. NVFP4 weight scales pad each matrix to whole 128-row atoms. When a gated I is not a multiple of 128, the gate rows of the canonical [up; gate] scale matrix start mid-atom, so CudnnFrostNvfp4Config.prepare_weights stores the FC1 block scales per branch. - Every SM120 source of every dtype is re-exported from cuDNN Frontend develop 886d7176. Kernel IDs and shortlists are unchanged, and the SM107 sources are byte-identical. - Exporter: SwiGLUStep's up clamp is now a slope-free relu(lower_clip=-7, upper_clip=7). cuDNN's RELU_FWD continues a slope below lower_clip, which Frost implements since cuDNN Frontend e0185b3b; the old relu(negative_slope=1, lower_clip=-7) spelling then leaks below -7. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 5.5 --- flashinfer/fused_moe/api.py | 35 +- .../fused_moe/backends/cudnn_frost/README.md | 93 +++-- .../backends/cudnn_frost/activations.py | 54 ++- .../bf16/cudnn_frost_selected_kernels.json | 122 ++++--- .../sm120_fc1_gated_normal_bf16_stg.py | 5 + .../sources/sm120_fc1_gated_swap_bf16_stg.py | 5 + .../sm120_fc1_non_gated_normal_bf16_stg.py | 5 + .../sm120_fc1_non_gated_swap_bf16_stg.py | 5 + .../sm120_fc2_non_gated_normal_bf16_stg.py | 5 + .../sm120_fc2_non_gated_swap_bf16_stg.py | 5 + .../mxfp8/cudnn_frost_selected_kernels.json | 220 +++++++----- .../sm120_fc1_gated_normal_bf16_stg.py | 13 +- .../sm120_fc1_gated_normal_quantized_stg.py | 13 +- .../sources/sm120_fc1_gated_swap_bf16_stg.py | 13 +- .../sm120_fc1_non_gated_normal_bf16_stg.py | 13 +- ...m120_fc1_non_gated_normal_quantized_stg.py | 13 +- .../sm120_fc1_non_gated_swap_bf16_stg.py | 13 +- .../sm120_fc2_non_gated_normal_bf16_stg.py | 13 +- .../sm120_fc2_non_gated_swap_bf16_stg.py | 13 +- .../cudnn_frost_selected_kernels.json | 180 ++++++---- .../sm120_fc1_gated_normal_bf16_stg.py | 13 +- .../sm120_fc1_gated_normal_quantized_stg.py | 13 +- .../sources/sm120_fc1_gated_swap_bf16_stg.py | 13 +- .../sm120_fc1_non_gated_normal_bf16_stg.py | 13 +- ...m120_fc1_non_gated_normal_quantized_stg.py | 13 +- .../sm120_fc1_non_gated_swap_bf16_stg.py | 13 +- .../sm120_fc2_non_gated_normal_bf16_stg.py | 13 +- .../sm120_fc2_non_gated_swap_bf16_stg.py | 13 +- .../nvfp4/cudnn_frost_selected_kernels.json | 121 +++++-- .../sm120_fc1_gated_normal_bf16_stg.py | 13 +- .../sources/sm120_fc1_gated_swap_bf16_stg.py | 13 +- .../sm120_fc1_non_gated_normal_bf16_stg.py | 13 +- .../sm120_fc1_non_gated_swap_bf16_stg.py | 13 +- .../sm120_fc2_non_gated_normal_bf16_stg.py | 13 +- .../sm120_fc2_non_gated_swap_bf16_stg.py | 13 +- .../backends/cudnn_frost/bf16/moe.py | 51 ++- .../backends/cudnn_frost/bf16/runtime.py | 19 +- .../backends/cudnn_frost/bf16/support.py | 5 +- .../fused_moe/backends/cudnn_frost/catalog.py | 13 +- .../backends/cudnn_frost/csrc/moe_bf16.cu | 88 +++-- .../backends/cudnn_frost/csrc/moe_mxfp8.cu | 104 +++--- .../cudnn_frost/csrc/moe_mxfp8_mxfp4.cu | 104 +++--- .../backends/cudnn_frost/csrc/moe_nvfp4.cu | 105 ++++-- .../backends/cudnn_frost/dimension_catalog.py | 75 ++++- .../fused_moe/backends/cudnn_frost/export.py | 54 ++- .../backends/cudnn_frost/mxfp8/moe.py | 33 +- .../backends/cudnn_frost/mxfp8/runtime.py | 31 +- .../backends/cudnn_frost/mxfp8_mxfp4/moe.py | 33 +- .../cudnn_frost/mxfp8_mxfp4/runtime.py | 31 +- .../backends/cudnn_frost/nvfp4/moe.py | 57 +++- .../backends/cudnn_frost/nvfp4/runtime.py | 31 +- .../backends/cudnn_frost/nvfp4/support.py | 5 +- .../fused_moe/backends/cudnn_frost/support.py | 13 +- flashinfer/fused_moe/prepare.py | 76 +++++ tests/moe/test_cudnn_frost_kernels.py | 70 ++-- tests/moe/test_cudnn_frost_swiglu_param.py | 318 ++++++++++++++++++ 56 files changed, 1959 insertions(+), 498 deletions(-) create mode 100644 tests/moe/test_cudnn_frost_swiglu_param.py diff --git a/flashinfer/fused_moe/api.py b/flashinfer/fused_moe/api.py index 66bf0a36943..a804ca330dc 100644 --- a/flashinfer/fused_moe/api.py +++ b/flashinfer/fused_moe/api.py @@ -2944,16 +2944,43 @@ def supported(cls, arch: int) -> bool: class CudnnFrostNvfp4Config: """Frost NVFP4 MoE on SM107 and SM120, with packed precomputed routing. - Register prepared weights under ``cudnn_frost_nvfp4``. The canonical - ``cutlass_nvfp4`` view is also accepted without conversion. H/I must be - multiples of 128. Pass the layer's quant config to prepare_activations. + Register prepared weights under ``cudnn_frost_nvfp4``. H/I must be + multiples of 128 on SM107 and of 64 on SM120 (e.g. gpt-oss, 2880). The + canonical ``cutlass_nvfp4`` view is also accepted without conversion, + except for a gated activation whose I is not a multiple of 128: its FC1 + block scales must be split per branch, which ``prepare_weights`` does. + Pass the layer's quant config to prepare_activations. """ @classmethod def supported(cls, arch: int) -> bool: return arch in (107, 120) - prepare_weights = staticmethod(CutlassNvfp4Config.prepare_weights) + @staticmethod + def prepare_weights( + w1_bf16, + w2_bf16, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: Optional[ActivationConfig] = None, + device=None, + ): + """Quantize canonical BF16 weights into the CUTLASS NVFP4 view, with + per-branch FC1 block scales where Frost needs them.""" + from .prepare import prepare_cudnn_frost_nvfp4_weights + + return prepare_cudnn_frost_nvfp4_weights( + w1_bf16, + w2_bf16, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + activation=activation, + device=device, + ) + prepare_activations = staticmethod(CutlassNvfp4Config.prepare_activations) diff --git a/flashinfer/fused_moe/backends/cudnn_frost/README.md b/flashinfer/fused_moe/backends/cudnn_frost/README.md index 4e58c85ce89..1bc39887024 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/README.md +++ b/flashinfer/fused_moe/backends/cudnn_frost/README.md @@ -7,10 +7,12 @@ SHA-256 manifest, with no precompiled `.o` files. Required cuDNN Frost device he are included in each generated source. The deployed process does not import `cudnn-frontend`, construct a cuDNN Frost graph, or run the cuDNN Frost graph compiler. -All packaged Tensor Core source families are generated from cuDNN Frontend -`0d72d59efd21401c7e13d606584fef2544e45934`, including ordinary FC1/FC2, -dedicated GATHER, and qualified NVFP4 COMBINE kernels. This revision is shared -by BF16, NVFP4, MXFP8 and MXFP8 × MXFP4 on their packaged SM107a/SM120a paths. +The packaged SM107a Tensor Core source families are generated from cuDNN +Frontend `0d72d59efd21401c7e13d606584fef2544e45934`, including ordinary +FC1/FC2, dedicated GATHER, and qualified NVFP4 COMBINE kernels; the SM120a +families from its descendant `886d7176` (see the SM120 section below). Each +revision is shared by BF16, NVFP4, MXFP8 and MXFP8 × MXFP4 on that +architecture's packaged paths. The catalogs retain their existing tile, cluster, activation and routing-mode coverage. Geometry constants are reconstructed from the producer's source; device-dependent pipeline depths and descriptor workspace sizes are resolved @@ -248,7 +250,9 @@ Per-expert `fc1_input_scale` and `fc2_input_scale` are runtime FP32 multipliers. FC1 applies its multiplier to the GEMM accumulators before activation, including both gate and up branches; FC2 applies its multiplier before BF16 output conversion. Bias, shared experts, expert parallelism, -non-default activation parameters and logits-based routing are unsupported. +non-default activation parameters (except SwiGLU's alpha/beta/limit on SM120, +see [Launch-time SwiGLU parameters](#launch-time-swiglu-parameters-sm120)) and +logits-based routing are unsupported. Automatic admission requires SM107a or SM120a, MXFP8 operands, BF16 finalized output, precomputed routing and legal artifacts for both stages. No @@ -259,7 +263,8 @@ the next measured profile; native executors use the actual input shape at launch MXFP8, NVFP4 and MXFP8 × MXFP4 shapes without a table entry (including token counts beyond its largest profile) use a bounded heuristic pool. H/I must be -positive multiples of 128, at most 2^20; `1 <= top_k <= E <= 1024`; +positive multiples of 128 (of 64 for NVFP4 on SM120), at most 2^20; +`1 <= top_k <= E <= 1024`; `1 <= T <= min(2^20, tune_max_num_tokens)`; expanded rows must satisfy the native int32 bound. Candidate selection checks each stage's artifact dimension and activation contracts; missing legal stages are reported during packing. @@ -388,7 +393,12 @@ templates for different architectures can coexist in each dtype directory. All ten default activation contracts are included: gated `SwiGLU`, `GeGLU`, `GeGLUTanh`, `SwiGLUStep`, `SiTU`; non-gated `Identity`, `ReLU`, `ReLU2`, `GELU`, `SiLU`. SiTU binds gate/linear scales 4/25 in the generated auxiliary -argument order; SwiGLUStep uses limit 7. FC1 fuses activation on FP32 +argument order; SwiGLUStep uses limit 7. SM120 additionally packages +`swiglu_param` (`--op grouped_gemm1_swiglu_param`), SwiGLU with launch-time +alpha/beta/limit. The exporter spells SwiGLUStep's up clamp as a slope-free +`relu(lower_clip=-7, upper_clip=7)`: cuDNN's RELU_FWD lets a slope continue +below `lower_clip`, which Frost implements since cuDNN Frontend `e0185b3b`, and +both packaged revisions are later than that. FC1 fuses activation on FP32 accumulators before the BF16 output conversion. The optional `scale` argument multiplies the activation output for every FC1, including non-gated variants. @@ -533,8 +543,15 @@ script's complete routed MoE measurements. `nvfp4/runtime.py`, `nvfp4/moe.py` and `csrc/moe_nvfp4.cu` provide the corresponding NVFP4 × NVFP4 pipeline and `cudnn_frost_nvfp4` automatic candidate. Both operands contain two E2M1 values per byte, with E4M3 block -scales per 16 logical K elements. Logical N and K must be divisible by 128; -packing halves the tensor's last storage dimension, not its logical K. +scales per 16 logical K elements. Logical N and K of the standalone grouped +API must be divisible by 128; packing halves the tensor's last storage +dimension, not its logical K. The MoE pipeline takes H and I in multiples of +64 on SM120 (whole 4-column scale words, e.g. gpt-oss's 2880) and of 128 on +SM107. Weight scales then pad each matrix to whole 128-row atoms: +`fc2_weight_block_scale` is `[E, round_up(H, 128), I/16]`, and when a gated I +is not a multiple of 128 the gate rows of the canonical `[up; gate]` matrix +start mid-atom, so `CudnnFrostNvfp4Config.prepare_weights` stores +`fc1_weight_block_scale` per branch as `[2, E, round_up(I, 128), H/16]`. The native pipeline gathers packed inputs, runs FC1/activation and intermediate NVFP4 quantization, runs FC2 with BF16 output, and reduces routed outputs in FP32 before the final BF16 conversion. @@ -723,11 +740,37 @@ checked against the architecture, geometry and activation in the manifest. Set `MoEConfig(activation=...)` using the existing typed activation API. Non-default scalar parameters and per-expert overrides are rejected rather -than silently ignored. In particular, selected SiTU uses gate scale 4 and -linear scale 25; SwiGLUStep uses limit 7. FC1 applies the activation to FP32 -accumulators before storing BF16. Gated activations fuse two grouped GEMMs; -non-gated activations use one. FC2 sources are shared across activations. -The standalone `cudnn_frost_grouped_gemm1_swiglu` API remains SwiGLU-specific. +than silently ignored, except SwiGLU's on SM120 (below). In particular, +selected SiTU uses gate scale 4 and linear scale 25; SwiGLUStep uses limit 7. +FC1 applies the activation to FP32 accumulators before storing BF16. Gated +activations fuse two grouped GEMMs; non-gated activations use one. FC2 sources +are shared across activations. The standalone +`cudnn_frost_grouped_gemm1_swiglu` API remains SwiGLU-specific. + +### Launch-time SwiGLU parameters (SM120) + +`SwiGLU(alpha, beta, limit)` with any non-default value runs the +`swiglu_param` sources, which read the parameters as FP32 scalar kernel +inputs, so one kernel serves every model: + +```text +g = min(gate, limit); u = clamp(up, -limit, limit) +out = g * sigmoid(alpha * g) * (u + beta) * scale +``` + +This is FlashInfer's SwiGLU contract (`tests/moe/utils.py`): e.g. +DeepSeek-V4/V4.1-Flash, its DSpark draft head and GLM-5.3-Flash use +`SwiGLU(limit=10)`; gpt-oss and MiniMax-M3 use +`SwiGLU(alpha=1.702, beta=1, limit=7)`. The default `SwiGLU()` keeps its +`swiglu` sources. `swiglu_param` is packaged for SM120 in every dtype, with +the same tiles, orientations and output variants as `swiglu`; on SM107 the +runner declines it with `NotImplementedError` instead of dropping the +parameters. Its launch tail places `swiglu_limit`, `swiglu_alpha`, +`swiglu_lower` and `swiglu_beta` (Frost's first-use order) before `scale`; +the native executors write `alpha, beta, limit, -limit` to scalar slots during +routing, so CUDA Graph replays keep them. `activations.py` maps an activation +to its source name (`activation_name`) and its launch values +(`launch_scalars`). The native executor carries the generated argument order, including SiTU's scalar inputs, for both normal and swap-AB layouts. Activation-specific source hashes @@ -1033,12 +1076,15 @@ activation of the topology expands through its exact source fragments. Kernel IDs start with `sm120_cudnn_frost_`, and catalog and shortlist records carry `arch: sm_120a`. -The current SM120 producer revision is listed at the start of this document. -It includes TMA L2 promotion on every SM120 template and sizes the packed -E2M1 side of a mixed pair at its 64-byte row in the SMEM footprint check, -as `device.py` does here. It also masks K-padding scale bytes in block-scale -mainloops and uses an unsigned ballot when finding the scheduler's winning -lane. Each host +The SM120 sources are exported from cuDNN Frontend `886d7176` (develop, +2026-11-04), a descendant of the SM107 revision listed at the start of this +document. It includes TMA L2 promotion on every SM120 template and sizes the +packed E2M1 side of a mixed pair at its 64-byte row in the SMEM footprint +check, as `device.py` does here. It also masks K-padding scale bytes in +block-scale mainloops, uses an unsigned ballot when finding the scheduler's +winning lane, clamps each MoE group to its routed rows, follows cuDNN's ReLU +slope/clip definition and propagates NaN through ReLU-family epilogues. Each +host zeroes its scheduler counter with a one-thread kernel before the PDL main launch, as the SM100 MoE hosts do, so the exported TVM-FFI function needs no caller-side reset. Like the SM107 sources, the schedulers take `G+1` explicit @@ -1053,8 +1099,11 @@ after the tile constants; `extract_template` keeps their decorators. The frozen source of a tile does not depend on E, N or K: exports of one tile at several model geometries render identical templates and parameters and coalesce into one logical kernel. Every SM120 kernel therefore uses the -catalog's generic contract (N and K multiples of 128, up to 1024 experts), as -on SM107. MXFP8 and MXFP8 × MXFP4 FC1 also +catalog's generic contract: N and K multiples of 32 (the kernels take whole +32-element blocks and mask the K tail), up to 1024 experts; SM107 kernels +need multiples of 128. Above the kernels, BF16 admits H/I in 32s and NVFP4 in +64s, while MXFP8 and MXFP8 × MXFP4 keep 128s: their canonical weight scales +are int32-packed per 128 K. MXFP8 and MXFP8 × MXFP4 FC1 also fuse the intermediate quantization (E4M3 data and segmented E8M0 scales, both through STG) in the normal orientation, for the tiles whose warp tile drains whole 32-column quantization blocks. cuDNN Frost cannot store packed FP4 on diff --git a/flashinfer/fused_moe/backends/cudnn_frost/activations.py b/flashinfer/fused_moe/backends/cudnn_frost/activations.py index 22fbe73d4f8..eb39ad4acee 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/activations.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/activations.py @@ -13,6 +13,7 @@ SwiGLUStep, ) +# Activations whose selected sources bake in the activation's default parameters. ACTIVATIONS = { "swiglu": SwiGLU, "geglu": GeGLU, @@ -25,19 +26,70 @@ "silu": SiLU, "identity": Identity, } -_DEFAULT_ACTIVATIONS = {name: cls() for name, cls in ACTIVATIONS.items()} +# Activations whose sources read their parameters as launch-time scalars, so one +# kernel serves every value. ``swiglu_param`` is SwiGLU with non-default +# ``alpha``, ``beta`` or ``limit``: ``g = min(gate, limit)``, +# ``u = clamp(up, -limit, limit)``, ``g * sigmoid(alpha * g) * (u + beta)``. +PARAMETERIZED_ACTIVATIONS = {"swiglu_param": SwiGLU} +# Every activation name a catalog record or exported op may carry. +KNOWN_ACTIVATIONS = {**ACTIVATIONS, **PARAMETERIZED_ACTIVATIONS} +# Launch-tail names of the swiglu_param scalars; the lower bound is -limit. +SWIGLU_PARAM_SCALARS = ("swiglu_alpha", "swiglu_beta", "swiglu_limit", "swiglu_lower") +_DEFAULT_ACTIVATIONS = {name: cls() for name, cls in KNOWN_ACTIVATIONS.items()} def activation_name(activation): for name, cls in ACTIVATIONS.items(): if type(activation) is cls and activation == _DEFAULT_ACTIVATIONS[name]: return name + if type(activation) is SwiGLU: + # Non-default alpha/beta/limit (e.g. DeepSeek-V4's limit 10, gpt-oss's + # alpha 1.702 / beta 1 / limit 7) run the scalar-parameterized sources. + return "swiglu_param" raise NotImplementedError( "cuDNN Frost selected sources require default activation parameters; " f"got {activation!r}" ) +def activation_scalars(activation) -> dict[str, float]: + """Launch-time scalar values (launch-tail name -> value) of an activation. + + Empty for activations whose parameters are compiled into their sources. + """ + if activation is None or activation_name(activation) != "swiglu_param": + return {} + limit = float(activation.limit) + return dict( + swiglu_alpha=float(activation.alpha), + swiglu_beta=float(activation.beta), + swiglu_limit=limit, + swiglu_lower=-limit, + ) + + +_FLOAT32_MAX = 3.4028234663852886e38 +_SCALAR_DEFAULTS = (1.0, 0.0, _FLOAT32_MAX, -_FLOAT32_MAX) + + +def launch_scalars(activation) -> tuple[float, float, float, float]: + """``(alpha, beta, limit, -limit)`` for the native MoE executors' scalar slots. + + Kernels without launch-time parameters never read them; they then hold + SwiGLU's defaults. + """ + values = activation_scalars(activation) + alpha, beta, limit, lower = ( + values.get(name, default) + for name, default in zip(SWIGLU_PARAM_SCALARS, _SCALAR_DEFAULTS, strict=True) + ) + return alpha, beta, limit, lower + + +# Scalar slots of executors whose kernels take no launch-time parameters. +DEFAULT_LAUNCH_SCALARS = launch_scalars(None) + + def contract_activation(contract): # Original schema-2 SwiGLU sources predate the typed activation key. value = contract.get("activation") diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/cudnn_frost_selected_kernels.json b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/cudnn_frost_selected_kernels.json index 672c8341762..ffbbecb25ef 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/cudnn_frost_selected_kernels.json +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/cudnn_frost_selected_kernels.json @@ -7,6 +7,14 @@ "swiglu_step", "situ" ], + "gated_param": [ + "swiglu", + "geglu", + "geglu_tanh", + "swiglu_step", + "situ", + "swiglu_param" + ], "non_gated": [ "gelu", "identity", @@ -5160,7 +5168,7 @@ ] }, { - "activation_group": "gated", + "activation_group": "gated_param", "arch": "sm_120a", "stage": "fc1", "variants": [ @@ -5194,9 +5202,9 @@ "56112e941b571de6", "7a4225debec6eddc", "2ed87fd309985059", - "4d401b076ec024cd", + "220136b3e0ca7cc7", "27d6d011a3008912", - "897db959e788d1d6", + "5092095203c53e38", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -5214,6 +5222,18 @@ "e3b0c44298fc1c14", "2f2e3279e3c767c7" ], + "swiglu_param": [ + "2754a26085d00492", + "3b18f23975a95a83", + "c2851a158d6a0bb3", + "8ce67ce32c873698", + "61e28c519656cc78", + "2135b8d3f20cfcda", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "b8cedbdf6cc2fc79" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -5238,13 +5258,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -5258,7 +5278,7 @@ "producer_template": "sm120_moe_grouped_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_gated_normal_bf16_stg.py", - "sha256": "4e565fa050212d00d821a4cff60896b5b9f157f28ebd15ce79a327d808f62905" + "sha256": "19d7c68d5cab2bee5d741418bec8d64110d74a62d34312613096106e93a8b19b" }, "store_policy": "stg", "tile_parameters": { @@ -5493,10 +5513,10 @@ "56112e941b571de6", "eb00231ff480d1e5", "2ed87fd309985059", - "4d401b076ec024cd", + "220136b3e0ca7cc7", "7bea654a536d6bf5", "27d6d011a3008912", - "897db959e788d1d6", + "5092095203c53e38", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -5515,6 +5535,19 @@ "e3b0c44298fc1c14", "d749578072d1c69f" ], + "swiglu_param": [ + "2754a26085d00492", + "d5f025c6d2c92c97", + "c2851a158d6a0bb3", + "8ce67ce32c873698", + "dc101805d6303043", + "61e28c519656cc78", + "2135b8d3f20cfcda", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "de25376457a0ccc7" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -5540,13 +5573,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -5560,7 +5593,7 @@ "producer_template": "sm120_moe_grouped_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_gated_swap_bf16_stg.py", - "sha256": "150ec2c84d47bb062158940f554246949cea0a20330a161ce6dbc7a52dc3a26e" + "sha256": "841da895ec587c3167ea0598c1155c383812e80d01ae724d2bada103128426a1" }, "store_policy": "stg", "tile_parameters": { @@ -5878,13 +5911,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -5898,7 +5931,7 @@ "producer_template": "sm120_moe_grouped_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_non_gated_normal_bf16_stg.py", - "sha256": "7b148982d0690ca014d970da9559a8dd7abac9a01c24fc62fedbe4c392982a24" + "sha256": "b1dcf660922c676acd8aec4ddaaefaa7a24360df6894f8f72de3a38e72319833" }, "store_policy": "stg", "tile_parameters": { @@ -6088,13 +6121,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -6108,7 +6141,7 @@ "producer_template": "sm120_moe_grouped_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_non_gated_swap_bf16_stg.py", - "sha256": "3d7078d835a61dce9c9851767addaa76662cf2225cc9fcc5b72fe29100e8c27f" + "sha256": "fbb5ed545c137449708b6c40751410d064c4e00eb1b88e0fe4111933df2b6ee7" }, "store_policy": "stg", "tile_parameters": { @@ -6308,13 +6341,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -6328,7 +6361,7 @@ "producer_template": "sm120_moe_grouped_matmul_fwd.py", "source": { "path": "sources/sm120_fc2_non_gated_normal_bf16_stg.py", - "sha256": "6c877a37dd9f6a1608cb3fd6c4834564cf6dc22855234655e4e85971ff3b452e" + "sha256": "5eab8f079a76c5c4fce79dfe4dd8ca3a6303214ff8324df0d6c4488ec0061ca4" }, "store_policy": "stg", "tile_parameters": { @@ -6442,13 +6475,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -6462,7 +6495,7 @@ "producer_template": "sm120_moe_grouped_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc2_non_gated_swap_bf16_stg.py", - "sha256": "b62542a22ea2b2c0e2c039438079cd68d1c1e9a326975bd0d4d9367a597c494c" + "sha256": "a68666dff063236bf3d65450d76b4a5a9c54f4f66cbdc7c0b57bafb20a82fc25" }, "store_policy": "stg", "tile_parameters": { @@ -6588,7 +6621,10 @@ "1e0f88de1db195e1": " _op_1 = (cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_1_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_1_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_2_a\n _c_3_a = (vec_f32).to(cutlass.Float32)\n _op_3 = cute.math.rcp(cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_4_a = (_op_2).to(cutlass.Float32)\n _c_4_b = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * _c_4_b\n _c_5_a = (vec_f32_1).to(cutlass.Float32)\n _op_5 = _c_5_a / (cutlass.full_like(_c_5_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = (cutlass.full_like(_c_6_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_6_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_6_a * cutlass.full_like(_c_6_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_6_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_7_a\n _c_8_a = (_op_4).to(cutlass.Float32)\n _c_8_b = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * _c_8_b\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * (cutlass.full_like(_c_9_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_9 = (_op_9).to(cutlass.BFloat16)\n _tap_0 = (_r_9).to(cutlass.BFloat16)\n", "201b9b3f6a65a887": " _c_1_b = (vec_f32_1).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "202cf2c641e8df8c": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_c_tap_0, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "2135b8d3f20cfcda": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_3_a = (_op_0).to(cutlass.Float32)\n _c_3_b = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (vec_f32_1).to(cutlass.Float32)\n _op_4 = cute.math.min(_c_4_a, (cutlass.full_like(_c_4_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_5_a = (_op_4).to(cutlass.Float32)\n _op_5 = cute.math.max(_c_5_a, (cutlass.full_like(_c_5_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a + (cutlass.full_like(_c_6_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_7_a = (_op_3).to(cutlass.Float32)\n _c_7_b = (_op_6).to(cutlass.Float32)\n _op_7 = _c_7_a * _c_7_b\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * (cutlass.full_like(_c_8_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_8 = (_op_8).to(cutlass.BFloat16)\n _tap_0 = (_r_8).to(cutlass.BFloat16)\n", "21ad07e05fe83f8d": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_c_tap_0, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "220136b3e0ca7cc7": " _th1 = _c_1_a\n _thq1 = _th1 * _th1\n _op_1 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th1) < cutlass.full_like(_th1, cutlass.Float32(0.125)), _th1 + _th1 * _thq1 * (cutlass.full_like(_thq1, cutlass.Float32(-0.3333333333333333)) + _thq1 * (cutlass.full_like(_thq1, cutlass.Float32(0.13333333333333333)) + _thq1 * (cutlass.full_like(_thq1, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th1, cutlass.Float32(1.0)) - cutlass.full_like(_th1, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th1 * cutlass.full_like(_th1, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th1, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_2_a\n _c_3_a = (vec_f32).to(cutlass.Float32)\n _op_3 = cute.math.rcp(cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_4_a = (_op_2).to(cutlass.Float32)\n _c_4_b = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * _c_4_b\n _c_5_a = (vec_f32_1).to(cutlass.Float32)\n _op_5 = _c_5_a / (cutlass.full_like(_c_5_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _th6 = _c_6_a\n _thq6 = _th6 * _th6\n _op_6 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th6) < cutlass.full_like(_th6, cutlass.Float32(0.125)), _th6 + _th6 * _thq6 * (cutlass.full_like(_thq6, cutlass.Float32(-0.3333333333333333)) + _thq6 * (cutlass.full_like(_thq6, cutlass.Float32(0.13333333333333333)) + _thq6 * (cutlass.full_like(_thq6, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th6, cutlass.Float32(1.0)) - cutlass.full_like(_th6, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th6 * cutlass.full_like(_th6, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th6, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_7_a\n _c_8_a = (_op_4).to(cutlass.Float32)\n _c_8_b = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * _c_8_b\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * (cutlass.full_like(_c_9_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_9 = (_op_9).to(cutlass.BFloat16)\n _tap_0 = (_r_9).to(cutlass.BFloat16)\n", + "2754a26085d00492": " swiglu_limit: cute.Tensor,\n swiglu_alpha: cute.Tensor,\n swiglu_lower: cute.Tensor,\n swiglu_beta: cute.Tensor,\n", "27d6d011a3008912": " _op_0 = _c_0_a / (cutlass.full_like(_c_0_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "28e39a716817853e": " _op_0 = _c_0_a * cute.math.rcp(cutlass.full_like(_c_0_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_0_a * cutlass.full_like(_c_0_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "2a58221a954782af": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", @@ -6597,14 +6633,16 @@ "2f2e3279e3c767c7": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_c_tap_0, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "30bcd14235a5b41f": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "39694a0c019c84b3": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _c_1_b = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "3b18f23975a95a83": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "3bc9e968a5bc8c1d": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "3f3f09f806845f87": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "44b6981accc03913": " _c_1_b = (vec_f32_1).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "487e7a9c66b50b8a": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "4d1c6e8eb77500a7": " _e_x0 = _c_0_a\n _e_erf0 = cute.math.erf(_e_x0 * cutlass.full_like(_e_x0, cutlass.Float32(0.7071067811865475)))\n _op_0 = cutlass.full_like(_e_x0, cutlass.Float32(0.5)) * _e_x0 * (cutlass.full_like(_e_x0, cutlass.Float32(1.0)) + _e_erf0)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", - "4d401b076ec024cd": " _op_1 = (cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_1_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_1_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_2_a\n _c_3_a = (vec_f32).to(cutlass.Float32)\n _op_3 = cute.math.rcp(cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_4_a = (_op_2).to(cutlass.Float32)\n _c_4_b = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * _c_4_b\n _c_5_a = (vec_f32_1).to(cutlass.Float32)\n _op_5 = _c_5_a / (cutlass.full_like(_c_5_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = (cutlass.full_like(_c_6_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_6_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_6_a * cutlass.full_like(_c_6_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_6_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_7_a\n _c_8_a = (_op_4).to(cutlass.Float32)\n _c_8_b = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * _c_8_b\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * (cutlass.full_like(_c_9_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_9 = (_op_9).to(cutlass.BFloat16)\n _tap_0 = (_r_9).to(cutlass.BFloat16)\n", "4d7126d684ddc92c": " _e_x0 = _c_0_a\n _e_erf0 = cute.math.erf(_e_x0 * cutlass.full_like(_e_x0, cutlass.Float32(0.7071067811865475)))\n _op_0 = cutlass.full_like(_e_x0, cutlass.Float32(0.5)) * _e_x0 * (cutlass.full_like(_e_x0, cutlass.Float32(1.0)) + _e_erf0)\n", + "5092095203c53e38": " _th1 = _c_1_a\n _thq1 = _th1 * _th1\n _op_1 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th1) < cutlass.full_like(_th1, cutlass.Float32(0.125)), _th1 + _th1 * _thq1 * (cutlass.full_like(_thq1, cutlass.Float32(-0.3333333333333333)) + _thq1 * (cutlass.full_like(_thq1, cutlass.Float32(0.13333333333333333)) + _thq1 * (cutlass.full_like(_thq1, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th1, cutlass.Float32(1.0)) - cutlass.full_like(_th1, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th1 * cutlass.full_like(_th1, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th1, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_2_a\n _c_3_a = (vec_f32).to(cutlass.Float32)\n _op_3 = cute.math.rcp(cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_4_a = (_op_2).to(cutlass.Float32)\n _c_4_b = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * _c_4_b\n _c_5_a = (vec_f32_1).to(cutlass.Float32)\n _op_5 = _c_5_a / (cutlass.full_like(_c_5_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _th6 = _c_6_a\n _thq6 = _th6 * _th6\n _op_6 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th6) < cutlass.full_like(_th6, cutlass.Float32(0.125)), _th6 + _th6 * _thq6 * (cutlass.full_like(_thq6, cutlass.Float32(-0.3333333333333333)) + _thq6 * (cutlass.full_like(_thq6, cutlass.Float32(0.13333333333333333)) + _thq6 * (cutlass.full_like(_thq6, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th6, cutlass.Float32(1.0)) - cutlass.full_like(_th6, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th6 * cutlass.full_like(_th6, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th6, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_7_a\n _c_8_a = (_op_4).to(cutlass.Float32)\n _c_8_b = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * _c_8_b\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * (cutlass.full_like(_c_9_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_9 = (_op_9).to(cutlass.BFloat16)\n _tap_0 = (_r_9).to(cutlass.BFloat16)\n", "56112e941b571de6": " gate_scale: cute.Tensor,\n linear_scale: cute.Tensor,\n", + "61e28c519656cc78": " _op_0 = cute.math.min(_c_0_a, (cutlass.full_like(_c_0_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", "651d2fba69ef8f2a": " _op_0 = _c_0_a * (cutlass.full_like(_c_0_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_0 = (_op_0).to(cutlass.BFloat16)\n _tap_0 = (_r_0).to(cutlass.BFloat16)\n", "665f4dbd94cb631b": " _op_0 = _c_0_a * cute.math.rcp(cutlass.full_like(_c_0_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_0_a * cutlass.full_like(_c_0_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "6b03984131318a83": " _g_x0 = _c_0_a\n _g_inner0 = cutlass.full_like(_g_x0, cutlass.Float32(0.7978845608028654)) * (_g_x0 + cutlass.full_like(_g_x0, cutlass.Float32(0.044715)) * _g_x0 * _g_x0 * _g_x0)\n _g_tanh0 = (cutlass.full_like(_g_inner0, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner0, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner0 * cutlass.full_like(_g_inner0, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner0, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_0 = cutlass.full_like(_g_x0, cutlass.Float32(0.5)) * _g_x0 * (cutlass.full_like(_g_x0, cutlass.Float32(1.0)) + _g_tanh0)\n", @@ -6620,7 +6658,7 @@ "7aef8cb6adfbb711": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_c_tap_0, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "7bea654a536d6bf5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", "87542e089a7bd42e": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", - "897db959e788d1d6": " _op_1 = (cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_1_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_1_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_2_a\n _c_3_a = (vec_f32).to(cutlass.Float32)\n _op_3 = cute.math.rcp(cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_4_a = (_op_2).to(cutlass.Float32)\n _c_4_b = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * _c_4_b\n _c_5_a = (vec_f32_1).to(cutlass.Float32)\n _op_5 = _c_5_a / (cutlass.full_like(_c_5_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = (cutlass.full_like(_c_6_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_6_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_6_a * cutlass.full_like(_c_6_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_6_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_7_a\n _c_8_a = (_op_4).to(cutlass.Float32)\n _c_8_b = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * _c_8_b\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * (cutlass.full_like(_c_9_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_9 = (_op_9).to(cutlass.BFloat16)\n _tap_0 = (_r_9).to(cutlass.BFloat16)\n", + "8ce67ce32c873698": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_3_a = (_op_0).to(cutlass.Float32)\n _c_3_b = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (vec_f32_1).to(cutlass.Float32)\n _op_4 = cute.math.min(_c_4_a, (cutlass.full_like(_c_4_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_5_a = (_op_4).to(cutlass.Float32)\n _op_5 = cute.math.max(_c_5_a, (cutlass.full_like(_c_5_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a + (cutlass.full_like(_c_6_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_7_a = (_op_3).to(cutlass.Float32)\n _c_7_b = (_op_6).to(cutlass.Float32)\n _op_7 = _c_7_a * _c_7_b\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * (cutlass.full_like(_c_8_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_8 = (_op_8).to(cutlass.BFloat16)\n _tap_0 = (_r_8).to(cutlass.BFloat16)\n", "8ee3a1f20e3f0fef": " _e_x0 = _c_0_a\n _e_erf0 = cute.math.erf(_e_x0 * cutlass.full_like(_e_x0, cutlass.Float32(0.7071067811865475)))\n _op_0 = cutlass.full_like(_e_x0, cutlass.Float32(0.5)) * _e_x0 * (cutlass.full_like(_e_x0, cutlass.Float32(1.0)) + _e_erf0)\n", "9a54610c2bff7920": " _e_x0 = _c_0_a\n _e_erf0 = cute.math.erf(_e_x0 * cutlass.full_like(_e_x0, cutlass.Float32(0.7071067811865475)))\n _op_0 = cutlass.full_like(_e_x0, cutlass.Float32(0.5)) * _e_x0 * (cutlass.full_like(_e_x0, cutlass.Float32(1.0)) + _e_erf0)\n", "9ca6c91821d8f7c2": " _op_0 = _c_0_a * cute.math.rcp(cutlass.full_like(_c_0_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_0_a * cutlass.full_like(_c_0_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _tap_0 = (_r_1).to(cutlass.BFloat16)\n", @@ -6628,19 +6666,26 @@ "b1038e2dba7c4f07": " _op_0 = _c_0_a * cute.math.rcp(cutlass.full_like(_c_0_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_0_a * cutlass.full_like(_c_0_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _tap_0 = (_r_1).to(cutlass.BFloat16)\n", "b130121bf40c4def": " _g_x0 = _c_0_a\n _g_inner0 = cutlass.full_like(_g_x0, cutlass.Float32(0.7978845608028654)) * (_g_x0 + cutlass.full_like(_g_x0, cutlass.Float32(0.044715)) * _g_x0 * _g_x0 * _g_x0)\n _g_tanh0 = (cutlass.full_like(_g_inner0, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner0, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner0 * cutlass.full_like(_g_inner0, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner0, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_0 = cutlass.full_like(_g_x0, cutlass.Float32(0.5)) * _g_x0 * (cutlass.full_like(_g_x0, cutlass.Float32(1.0)) + _g_tanh0)\n", "b1e4ded98b0b6add": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _tap_0 = (_r_1).to(cutlass.BFloat16)\n", + "b8cedbdf6cc2fc79": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_c_tap_0, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "b932e17e4b585343": " swiglu_limit,\n swiglu_alpha,\n swiglu_lower,\n swiglu_beta,\n", "b9ea48e8214972b2": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _tap_0 = (_r_1).to(cutlass.BFloat16)\n", "bf848eb4a38a4cb4": " _op_1 = (cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_1_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_1_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_2_a\n _c_3_a = (vec_f32).to(cutlass.Float32)\n _op_3 = cute.math.rcp(cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_4_a = (_op_2).to(cutlass.Float32)\n _c_4_b = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * _c_4_b\n _c_5_a = (vec_f32_1).to(cutlass.Float32)\n _op_5 = _c_5_a / (cutlass.full_like(_c_5_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = (cutlass.full_like(_c_6_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_6_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_6_a * cutlass.full_like(_c_6_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_6_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_7_a\n _c_8_a = (_op_4).to(cutlass.Float32)\n _c_8_b = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a * _c_8_b\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * (cutlass.full_like(_c_9_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_9 = (_op_9).to(cutlass.BFloat16)\n vec_out = (_r_9).to(cutlass.BFloat16)\n", "c19027e18fa4927f": " _e_x0 = _c_0_a\n _e_erf0 = cute.math.erf(_e_x0 * cutlass.full_like(_e_x0, cutlass.Float32(0.7071067811865475)))\n _op_0 = cutlass.full_like(_e_x0, cutlass.Float32(0.5)) * _e_x0 * (cutlass.full_like(_e_x0, cutlass.Float32(1.0)) + _e_erf0)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _tap_0 = (_r_1).to(cutlass.BFloat16)\n", "c1fcf91531de08a1": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _c_1_b = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "c2851a158d6a0bb3": " _op_0 = cute.math.min(_c_0_a, (cutlass.full_like(_c_0_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", "c8882d1b1215864f": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _c_1_b = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "d0bf38bdc7f3c829": " _c_1_b = (vec_f32_1).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "d5f025c6d2c92c97": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "d749578072d1c69f": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_c_tap_0, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "d74c1b876c2737f8": " fake_gate_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_linear_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "da102211af750a6d": " _op_0 = _c_0_a * (cutlass.full_like(_c_0_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_0 = (_op_0).to(cutlass.BFloat16)\n _tap_0 = (_r_0).to(cutlass.BFloat16)\n", + "dc101805d6303043": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "dc63a189abb0282b": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_c_tap_0, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "de009f30c7784f96": " _op_0 = cute.math.max(_c_0_a, cutlass.full_like(_c_0_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_1_a = (_op_0).to(cutlass.Float32)\n _c_1_b = (_op_0).to(cutlass.Float32)\n _op_1 = _c_1_a * _c_1_b\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "de25376457a0ccc7": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_c_tap_0, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "e17d1851605245f0": " _r1 = _c_1_a\n _rk1 = cute.math.max(_r1, cutlass.full_like(_r1, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r1, cutlass.Float32(1.0)) * cute.math.min(_r1, cutlass.full_like(_r1, cutlass.Float32(0.0)), propagate_nan=True)\n _ru1 = cute.math.min(_rk1, cutlass.full_like(_r1, cutlass.Float32(7.0)), propagate_nan=True)\n _op_1 = _ru1\n _c_2_a = (vec_f32_1).to(cutlass.Float32)\n _r2 = _c_2_a\n _rl2 = cute.math.max(_r2, cutlass.full_like(_r2, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru2 = cute.math.min(_rl2, cutlass.full_like(_r2, cutlass.Float32(7.0)), propagate_nan=True)\n _op_2 = _ru2\n _c_3_a = (_op_1).to(cutlass.Float32)\n _c_3_b = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _tap_0 = (_r_4).to(cutlass.BFloat16)\n", "e3b0c44298fc1c14": "", + "e75461c9caab2535": " fake_swiglu_limit = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_alpha = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_lower = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_beta = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "e82fa158517f70a0": " _r1 = _c_1_a\n _rk1 = cute.math.max(_r1, cutlass.full_like(_r1, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r1, cutlass.Float32(1.0)) * cute.math.min(_r1, cutlass.full_like(_r1, cutlass.Float32(0.0)), propagate_nan=True)\n _ru1 = cute.math.min(_rk1, cutlass.full_like(_r1, cutlass.Float32(7.0)), propagate_nan=True)\n _op_1 = _ru1\n _c_2_a = (vec_f32_1).to(cutlass.Float32)\n _r2 = _c_2_a\n _rl2 = cute.math.max(_r2, cutlass.full_like(_r2, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru2 = cute.math.min(_rl2, cutlass.full_like(_r2, cutlass.Float32(7.0)), propagate_nan=True)\n _op_2 = _ru2\n _c_3_a = (_op_1).to(cutlass.Float32)\n _c_3_b = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _tap_0 = (_r_4).to(cutlass.BFloat16)\n", "eb00231ff480d1e5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", "eb71869b46342011": " _op_0 = _c_0_a * cute.math.rcp(cutlass.full_like(_c_0_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_0_a * cutlass.full_like(_c_0_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", @@ -7443,7 +7488,8 @@ "producer": { "name": "cudnn_frost", "revisions": [ - "0d72d59efd21401c7e13d606584fef2544e45934" + "0d72d59efd21401c7e13d606584fef2544e45934", + "886d7176dbe724a05f0c3955c0a87abb326afd59" ] }, "schema_version": 4, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_normal_bf16_stg.py index 82d969a297d..cb425f19e17 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_normal_bf16_stg.py @@ -460,6 +460,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile count cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n prefix_tiles = my_tiles for delta in (1, 2, 4, 8, 16): diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_swap_bf16_stg.py index 37291cfa17e..7e5c632bc8c 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_gated_swap_bf16_stg.py @@ -458,6 +458,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile count cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m prefix_tiles = my_tiles for delta in (1, 2, 4, 8, 16): diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_normal_bf16_stg.py index dc9cc76ce2c..f2e85acef24 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_normal_bf16_stg.py @@ -458,6 +458,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile count cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n prefix_tiles = my_tiles for delta in (1, 2, 4, 8, 16): diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_swap_bf16_stg.py index 0404f559446..45100204b34 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc1_non_gated_swap_bf16_stg.py @@ -456,6 +456,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile count cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m prefix_tiles = my_tiles for delta in (1, 2, 4, 8, 16): diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_normal_bf16_stg.py index 0ce986d48d1..f822003534c 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_normal_bf16_stg.py @@ -457,6 +457,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile count cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n prefix_tiles = my_tiles for delta in (1, 2, 4, 8, 16): diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_swap_bf16_stg.py index 62810927edd..0ac3e7e1b88 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/bf16/sources/sm120_fc2_non_gated_swap_bf16_stg.py @@ -455,6 +455,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile count cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m prefix_tiles = my_tiles for delta in (1, 2, 4, 8, 16): diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/cudnn_frost_selected_kernels.json b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/cudnn_frost_selected_kernels.json index fc06efea8ce..1285e2d1c63 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/cudnn_frost_selected_kernels.json +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/cudnn_frost_selected_kernels.json @@ -7,6 +7,14 @@ "swiglu_step", "situ" ], + "gated_param": [ + "swiglu", + "geglu", + "geglu_tanh", + "swiglu_step", + "situ", + "swiglu_param" + ], "non_gated": [ "gelu", "identity", @@ -7896,7 +7904,7 @@ ] }, { - "activation_group": "gated", + "activation_group": "gated_param", "arch": "sm_120a", "stage": "fc1", "variants": [ @@ -7930,9 +7938,9 @@ "56112e941b571de6", "7a4225debec6eddc", "41dc5fce5f0554b1", - "93767dfc83987731", + "41fd5323feb58adf", "29578b1cccc93958", - "d0718f1e440ae3aa", + "6f88d1729ff4d53e", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -7950,6 +7958,18 @@ "e3b0c44298fc1c14", "4328fc11f5cb9a36" ], + "swiglu_param": [ + "2754a26085d00492", + "3b18f23975a95a83", + "6c203c815404dd02", + "6ed3d2a160090752", + "b6f78000c3b8d982", + "f15d042ef49953d5", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "52d76c29c152ab54" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -7978,13 +7998,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8001,7 +8021,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_gated_normal_bf16_stg.py", - "sha256": "ad6a58f2068901cb60e1741627094c43e0f4f9e428c9110501a1cda45ad274fd" + "sha256": "c014b694785a06741a3a05cc53553b87612f81fa702305e4cd27b55b34af6896" }, "store_policy": "stg", "tile_parameters": { @@ -8011,6 +8031,7 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "4", + "mma_size_n": "1", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8023,7 +8044,6 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "2", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8048,7 +8068,6 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "4", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8065,6 +8084,7 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "2", + "mma_size_n": "1", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8076,7 +8096,6 @@ "cta_tile_mnk": "(32, 128, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8089,7 +8108,6 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "2", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8121,6 +8139,7 @@ "cta_tile_mnk": "(32, 64, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", + "mma_size_n": "1", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8135,7 +8154,6 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "2", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8160,7 +8178,6 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "4", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8192,6 +8209,7 @@ "cta_tile_mnk": "(64, 32, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", + "mma_size_n": "1", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8206,6 +8224,7 @@ "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "2", + "mma_size_n": "1", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8219,7 +8238,6 @@ "cta_tile_mnk": "(64, 64, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", - "mma_size_n": "2", "num_b_operands": "2", "num_gemms": "2", "num_sfb_operands": "2", @@ -8260,9 +8278,9 @@ "56112e941b571de6", "7a4225debec6eddc", "41dc5fce5f0554b1", - "425e3214fd5ae955", + "3eed40bd0e026d1e", "29578b1cccc93958", - "ffca1f68b5a79844", + "7a376b4eefc70b3c", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8280,6 +8298,18 @@ "e3b0c44298fc1c14", "b5d0f3bef19393b9" ], + "swiglu_param": [ + "2754a26085d00492", + "3b18f23975a95a83", + "6c203c815404dd02", + "0494fa69c080922a", + "b6f78000c3b8d982", + "066608727c187eb0", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "567045b85e6e674e" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -8309,13 +8339,13 @@ "intermediate_dtype": "bfloat16", "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "float8_e4m3fn", "output_scale_layout": "segmented_F8_128x4", @@ -8333,7 +8363,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_gated_normal_quantized_stg.py", - "sha256": "8393aa0eff0b74b33ca815ca823329a71d1eb55d80f65363deabbbe925b24282" + "sha256": "1900f462708738ce1463a2e41692220c6d039a1adc5a8d4eb9dea69b1500268d" }, "store_policy": "stg", "tile_parameters": { @@ -8458,10 +8488,10 @@ "56112e941b571de6", "eb00231ff480d1e5", "41dc5fce5f0554b1", - "93767dfc83987731", + "41fd5323feb58adf", "7bea654a536d6bf5", "29578b1cccc93958", - "d0718f1e440ae3aa", + "6f88d1729ff4d53e", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8480,6 +8510,19 @@ "e3b0c44298fc1c14", "a2ba14ff465ff35e" ], + "swiglu_param": [ + "2754a26085d00492", + "d5f025c6d2c92c97", + "6c203c815404dd02", + "6ed3d2a160090752", + "dc101805d6303043", + "b6f78000c3b8d982", + "f15d042ef49953d5", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "6f5ddd4860594594" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -8509,13 +8552,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8532,7 +8575,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_gated_swap_bf16_stg.py", - "sha256": "501d7f1e8e023fc963ab5cbc20ab3f8bc7fbe9c9cba4121b6ec5d082fd3d26af" + "sha256": "ec44c63520928119437d02ee662eb88fd3774ada1a3378779cadcae21d4a6648" }, "store_policy": "stg", "tile_parameters": { @@ -8542,6 +8585,7 @@ "epi_chunk_elems": "1", "gemm_a_idx": "(0, 1)", "gemm_b_idx": "(0, 0)", + "mma_size_n": "1", "num_a_operands": "2", "num_gemms": "2", "num_sfa_operands": "2", @@ -8556,7 +8600,6 @@ "epi_chunk_elems": "1", "gemm_a_idx": "(0, 1)", "gemm_b_idx": "(0, 0)", - "mma_size_n": "2", "num_a_operands": "2", "num_gemms": "2", "num_sfa_operands": "2", @@ -8571,6 +8614,7 @@ "epi_chunk_elems": "1", "gemm_a_idx": "(0, 1)", "gemm_b_idx": "(0, 0)", + "mma_size_n": "1", "num_a_operands": "2", "num_gemms": "2", "num_sfa_operands": "2", @@ -8588,7 +8632,6 @@ "gemm_a_idx": "(0, 1)", "gemm_b_idx": "(0, 0)", "mma_size_m": "2", - "mma_size_n": "2", "num_a_operands": "2", "num_gemms": "2", "num_sfa_operands": "2", @@ -8618,6 +8661,7 @@ "epi_chunk_elems": "1", "gemm_a_idx": "(0, 1)", "gemm_b_idx": "(0, 0)", + "mma_size_n": "1", "num_a_operands": "2", "num_gemms": "2", "num_sfa_operands": "2", @@ -8634,7 +8678,6 @@ "epi_chunk_elems": "1", "gemm_a_idx": "(0, 1)", "gemm_b_idx": "(0, 0)", - "mma_size_n": "2", "num_a_operands": "2", "num_gemms": "2", "num_sfa_operands": "2", @@ -8692,13 +8735,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8715,16 +8758,16 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_non_gated_normal_bf16_stg.py", - "sha256": "fb7ed918e296b9a97e5b1f24bf8187372f0ba9d7d0175905c46253f28cfeacbb" + "sha256": "2bfe5c522fb3bcd30bba3d3334d1f407f35874c055f71eace499fbc4424ebbc3" }, "store_policy": "stg", "tile_parameters": { "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps1x8": { - "mma_size_m": "8" + "mma_size_m": "8", + "mma_size_n": "1" }, "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps2x4": { - "mma_size_m": "4", - "mma_size_n": "2" + "mma_size_m": "4" }, "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps4x2": { "mma_size_m": "2", @@ -8737,13 +8780,13 @@ "cgrp_tile_mnk": "(32, 128, 128)", "cta_tile_mnk": "(32, 128, 128)", "mma_size_m": "2", + "mma_size_n": "1", "sA_packed_elems": "4096", "sA_tma_bytes": "4096" }, "CONFIG_sm120_32x128x128_16x16x32_cluster1x1_warps2x4": { "cgrp_tile_mnk": "(32, 128, 128)", "cta_tile_mnk": "(32, 128, 128)", - "mma_size_n": "2", "sA_packed_elems": "4096", "sA_tma_bytes": "4096" }, @@ -8751,7 +8794,6 @@ "cgrp_tile_mnk": "(32, 256, 128)", "cta_tile_mnk": "(32, 256, 128)", "mma_size_m": "2", - "mma_size_n": "2", "sA_packed_elems": "4096", "sA_tma_bytes": "4096", "sB_packed_elems": "32768", @@ -8773,6 +8815,7 @@ "CONFIG_sm120_32x64x128_16x16x32_cluster1x1_warps2x4": { "cgrp_tile_mnk": "(32, 64, 128)", "cta_tile_mnk": "(32, 64, 128)", + "mma_size_n": "1", "sA_packed_elems": "4096", "sA_tma_bytes": "4096", "sB_packed_elems": "8192", @@ -8782,6 +8825,7 @@ "cgrp_tile_mnk": "(64, 128, 128)", "cta_tile_mnk": "(64, 128, 128)", "mma_size_m": "4", + "mma_size_n": "1", "sA_packed_elems": "8192", "sA_tma_bytes": "8192" }, @@ -8789,7 +8833,6 @@ "cgrp_tile_mnk": "(64, 128, 128)", "cta_tile_mnk": "(64, 128, 128)", "mma_size_m": "2", - "mma_size_n": "2", "sA_packed_elems": "8192", "sA_tma_bytes": "8192" }, @@ -8804,6 +8847,7 @@ "cgrp_tile_mnk": "(64, 64, 128)", "cta_tile_mnk": "(64, 64, 128)", "mma_size_m": "2", + "mma_size_n": "1", "sA_packed_elems": "8192", "sA_tma_bytes": "8192", "sB_packed_elems": "8192", @@ -8812,7 +8856,6 @@ "CONFIG_sm120_64x64x128_16x16x32_cluster1x1_warps4x2": { "cgrp_tile_mnk": "(64, 64, 128)", "cta_tile_mnk": "(64, 64, 128)", - "mma_size_n": "2", "sA_packed_elems": "8192", "sA_tma_bytes": "8192", "sB_packed_elems": "8192", @@ -8859,13 +8902,13 @@ "intermediate_dtype": "bfloat16", "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "float8_e4m3fn", "output_scale_layout": "segmented_F8_128x4", @@ -8883,7 +8926,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_non_gated_normal_quantized_stg.py", - "sha256": "95d06c2e341c60431d799e04c78fea7ce6445a8cdb6e47912df0440c146d49bf" + "sha256": "d4d9ef0215dc83af858f1cd2a40dbff0a963476e0092e22d03705c71bb48a84c" }, "store_policy": "stg", "tile_parameters": { @@ -9034,13 +9077,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -9057,14 +9100,13 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_non_gated_swap_bf16_stg.py", - "sha256": "1089202c813f48fba32279c60690dd047ce056d9d24328e312cd2983ce3706b8" + "sha256": "1deb85b231a99b57f7577b6c8072a73a5c0d3039c34b09bbe37ec6b5d0873ba9" }, "store_policy": "stg", "tile_parameters": { "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps2x4_swapAB": { "epi_chunk_elems": "1", "mma_size_m": "4", - "mma_size_n": "2", "vec_bytes_epi": "2", "vec_bytes_tap_0": "2" }, @@ -9079,6 +9121,7 @@ "cgrp_tile_mnk": "(128, 16, 128)", "cta_tile_mnk": "(128, 16, 128)", "epi_chunk_elems": "1", + "mma_size_n": "1", "sB_packed_elems": "2048", "sB_tma_bytes": "2048", "vec_bytes_epi": "2", @@ -9089,6 +9132,7 @@ "cta_tile_mnk": "(128, 32, 128)", "epi_chunk_elems": "1", "mma_size_m": "2", + "mma_size_n": "1", "sB_packed_elems": "4096", "sB_tma_bytes": "4096", "vec_bytes_epi": "2", @@ -9098,7 +9142,6 @@ "cgrp_tile_mnk": "(128, 32, 128)", "cta_tile_mnk": "(128, 32, 128)", "epi_chunk_elems": "1", - "mma_size_n": "2", "sB_packed_elems": "4096", "sB_tma_bytes": "4096", "vec_bytes_epi": "2", @@ -9109,6 +9152,7 @@ "cta_tile_mnk": "(256, 16, 128)", "epi_chunk_elems": "1", "mma_size_m": "2", + "mma_size_n": "1", "sA_packed_elems": "32768", "sA_tma_bytes": "32768", "sB_packed_elems": "2048", @@ -9123,7 +9167,6 @@ "cta_tile_mnk": "(256, 32, 128)", "epi_chunk_elems": "1", "mma_size_m": "2", - "mma_size_n": "2", "sA_packed_elems": "32768", "sA_tma_bytes": "32768", "sB_packed_elems": "4096", @@ -9137,6 +9180,7 @@ "cgrp_tile_mnk": "(64, 32, 128)", "cta_tile_mnk": "(64, 32, 128)", "epi_chunk_elems": "1", + "mma_size_n": "1", "sA_packed_elems": "8192", "sA_tma_bytes": "8192", "sB_packed_elems": "4096", @@ -9172,13 +9216,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -9195,16 +9239,16 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc2_non_gated_normal_bf16_stg.py", - "sha256": "b490619653bcc5b1ce091defedc284714ad4031ee597e86a4ec53c9c1774acac" + "sha256": "e36235a8827b4f86d6c257542d18ed7f3746c4ec64c596dbbc3baf7cf35e7724" }, "store_policy": "stg", "tile_parameters": { "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps1x8": { - "mma_size_m": "8" + "mma_size_m": "8", + "mma_size_n": "1" }, "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps2x4": { - "mma_size_m": "4", - "mma_size_n": "2" + "mma_size_m": "4" }, "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps4x2": { "mma_size_m": "2", @@ -9216,7 +9260,6 @@ "CONFIG_sm120_32x128x128_16x16x32_cluster1x1_warps2x4": { "cgrp_tile_mnk": "(32, 128, 128)", "cta_tile_mnk": "(32, 128, 128)", - "mma_size_n": "2", "sA_packed_elems": "4096", "sA_tma_bytes": "4096" }, @@ -9224,7 +9267,6 @@ "cgrp_tile_mnk": "(32, 256, 128)", "cta_tile_mnk": "(32, 256, 128)", "mma_size_m": "2", - "mma_size_n": "2", "sA_packed_elems": "4096", "sA_tma_bytes": "4096", "sB_packed_elems": "32768", @@ -9246,6 +9288,7 @@ "CONFIG_sm120_32x64x128_16x16x32_cluster1x1_warps2x4": { "cgrp_tile_mnk": "(32, 64, 128)", "cta_tile_mnk": "(32, 64, 128)", + "mma_size_n": "1", "sA_packed_elems": "4096", "sA_tma_bytes": "4096", "sB_packed_elems": "8192", @@ -9255,7 +9298,6 @@ "cgrp_tile_mnk": "(64, 128, 128)", "cta_tile_mnk": "(64, 128, 128)", "mma_size_m": "2", - "mma_size_n": "2", "sA_packed_elems": "8192", "sA_tma_bytes": "8192" }, @@ -9269,6 +9311,7 @@ "CONFIG_sm120_64x32x128_16x16x32_cluster1x1_warps4x2": { "cgrp_tile_mnk": "(64, 32, 128)", "cta_tile_mnk": "(64, 32, 128)", + "mma_size_n": "1", "sA_packed_elems": "8192", "sA_tma_bytes": "8192", "sB_packed_elems": "4096", @@ -9295,13 +9338,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -9318,14 +9361,13 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc2_non_gated_swap_bf16_stg.py", - "sha256": "ca1254f0f7dca69e0eabd98badb311df951bddcb70023b68cfb539a9bde7063b" + "sha256": "799a25e845e62f32c50810856cbe29535c3e86e50c490b3f5a93539404cb466b" }, "store_policy": "stg", "tile_parameters": { "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps2x4_swapAB": { "epi_chunk_elems": "1", "mma_size_m": "4", - "mma_size_n": "2", "vec_bytes_epi": "2", "vec_bytes_tap_0": "2" }, @@ -9346,6 +9388,7 @@ "cgrp_tile_mnk": "(128, 16, 128)", "cta_tile_mnk": "(128, 16, 128)", "epi_chunk_elems": "1", + "mma_size_n": "1", "sB_packed_elems": "2048", "sB_tma_bytes": "2048", "vec_bytes_epi": "2", @@ -9356,6 +9399,7 @@ "cta_tile_mnk": "(256, 16, 128)", "epi_chunk_elems": "1", "mma_size_m": "2", + "mma_size_n": "1", "sA_packed_elems": "32768", "sA_tma_bytes": "32768", "sB_packed_elems": "2048", @@ -9369,6 +9413,7 @@ "cgrp_tile_mnk": "(64, 32, 128)", "cta_tile_mnk": "(64, 32, 128)", "epi_chunk_elems": "1", + "mma_size_n": "1", "sA_packed_elems": "8192", "sA_tma_bytes": "8192", "sB_packed_elems": "4096", @@ -9385,7 +9430,9 @@ "00c5ee6c27d970e0": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "015dae03e6eec8cf": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "047483e0c9f06583": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", + "0494fa69c080922a": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _q0_src = (_r_10).to(cutlass.Float32)\n", "0560352581e4245f": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", + "066608727c187eb0": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _q0_src = (_r_10).to(cutlass.Float32)\n", "097c496041ff0970": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _tap_0 = (_r_3).to(cutlass.BFloat16)\n", "09f974d1874b91c5": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", "0b80a44ad5fe8c55": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", @@ -9397,6 +9444,7 @@ "1dfc658c5c45d2e7": " gate_scale,\n linear_scale,\n", "1ffaeb97f374dcf1": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "22ca0d8f6cd1c1b0": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", + "2754a26085d00492": " swiglu_limit: cute.Tensor,\n swiglu_alpha: cute.Tensor,\n swiglu_lower: cute.Tensor,\n swiglu_beta: cute.Tensor,\n", "29578b1cccc93958": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "2c22943ab1bd98cf": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "2f39a5d044e06acf": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", @@ -9405,8 +9453,10 @@ "330cb022858093b7": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", "33aa3fa477b45f2f": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n vec_out = (_r_11).to(cutlass.BFloat16)\n", "378b021ec68df0b8": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", + "3b18f23975a95a83": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", + "3eed40bd0e026d1e": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "41dc5fce5f0554b1": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", - "425e3214fd5ae955": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", + "41fd5323feb58adf": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "430a595659be8727": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", "4328fc11f5cb9a36": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "458018ab4b9c29cd": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", @@ -9414,7 +9464,9 @@ "4babc9e5e7349459": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "50ed9889665429ca": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", "5265cd619c815f29": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "52d76c29c152ab54": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "56112e941b571de6": " gate_scale: cute.Tensor,\n linear_scale: cute.Tensor,\n", + "567045b85e6e674e": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "59195a7258d63f15": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", "5a54255f3f367e7e": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "5a8b7730d3a750c0": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", @@ -9422,9 +9474,14 @@ "61f8f91878afa3b0": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "663a4532415b8373": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "66ed7d44845792bf": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", + "6c203c815404dd02": " _op_2 = cute.math.min(_c_2_a, (cutlass.full_like(_c_2_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", + "6ed3d2a160090752": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _tap_0 = (_r_10).to(cutlass.BFloat16)\n", + "6f5ddd4860594594": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "6f88d1729ff4d53e": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "7023effb53be67be": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "7928eb53a5feaeef": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", "79bdcfbd725e3cdd": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", + "7a376b4eefc70b3c": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "7a4225debec6eddc": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", "7b83af455c81d304": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "7bea654a536d6bf5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", @@ -9443,7 +9500,6 @@ "90b5af501e75575f": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "90c60ce1df6c64fb": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n vec_out = (_r_3).to(cutlass.BFloat16)\n", "91ac55be27ccb00f": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", - "93767dfc83987731": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "9399417aa8da2382": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "941640545da1616a": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n vec_out = (_r_4).to(cutlass.BFloat16)\n", "949c2cdeadcfe453": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _tap_0 = (_r_4).to(cutlass.BFloat16)\n", @@ -9462,6 +9518,8 @@ "b47ef4c40992b214": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "b5d0f3bef19393b9": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "b6b4da9f62018b7e": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", + "b6f78000c3b8d982": " _op_2 = cute.math.min(_c_2_a, (cutlass.full_like(_c_2_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", + "b932e17e4b585343": " swiglu_limit,\n swiglu_alpha,\n swiglu_lower,\n swiglu_beta,\n", "bafb9857135fcddb": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "bcb449ff12f30cc5": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _q0_src = (_r_4).to(cutlass.Float32)\n", "bda8354eac44ce9a": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", @@ -9483,15 +9541,18 @@ "d41b4a8d60116de9": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "d43d926df1c8e8ed": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "d446e6dcb27a27f0": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", + "d5f025c6d2c92c97": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "d6aa0440273ef725": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", "d74c1b876c2737f8": " fake_gate_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_linear_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "d8077193e3e6cac7": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "d8b53fe287d3344e": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", + "dc101805d6303043": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "e3159eb3d0eec96f": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "e3b0c44298fc1c14": "", "e48523ab7fbb239b": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n vec_out = (_r_11).to(cutlass.BFloat16)\n", "e493bd216c43ecc2": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "e4ab4f7e50c8c2ee": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "e75461c9caab2535": " fake_swiglu_limit = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_alpha = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_lower = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_beta = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "e84fea998f6de219": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "ea8c14561c6061cf": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "eb00231ff480d1e5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", @@ -9501,11 +9562,11 @@ "ef1a34ae80069ad6": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "f0a06d037b34e7e2": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "f0cd44c3d5366545": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", + "f15d042ef49953d5": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _tap_0 = (_r_10).to(cutlass.BFloat16)\n", "f174b87147f89859": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", "f71332286643d5dc": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "fc9346af90b1bc7e": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "ff9d405c8dbd6211": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", - "ffca1f68b5a79844": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "ffde0b8549b9703a": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n" }, "parameter_bases": { @@ -10440,7 +10501,7 @@ "mma_scale_vec_size": "'X1'", "mma_sf_format": "'UE8M0'", "mma_size_m": "1", - "mma_size_n": "1", + "mma_size_n": "2", "moe_combine": "False", "moe_gather": "False", "moe_scatter": "False", @@ -10711,7 +10772,8 @@ "producer": { "name": "cudnn_frost", "revisions": [ - "0d72d59efd21401c7e13d606584fef2544e45934" + "0d72d59efd21401c7e13d606584fef2544e45934", + "886d7176dbe724a05f0c3955c0a87abb326afd59" ] }, "schema_version": 4, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_bf16_stg.py index 286a44bcc84..e0cc313c498 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_bf16_stg.py @@ -577,6 +577,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -606,6 +610,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -629,7 +638,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_quantized_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_quantized_stg.py index f9452faf423..f506a35dd11 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_quantized_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_normal_quantized_stg.py @@ -615,6 +615,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -644,6 +648,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -667,7 +676,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_swap_bf16_stg.py index ad473b5db72..ea330e431ac 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_gated_swap_bf16_stg.py @@ -576,6 +576,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -605,6 +609,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -628,7 +637,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_bf16_stg.py index 6a23b137690..07199d2691b 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_bf16_stg.py @@ -573,6 +573,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -602,6 +606,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -625,7 +634,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_quantized_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_quantized_stg.py index 5cca44e768a..1a34b9fd932 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_quantized_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_normal_quantized_stg.py @@ -611,6 +611,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -640,6 +644,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -663,7 +672,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_swap_bf16_stg.py index d3efbd69e44..5bf392a672a 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc1_non_gated_swap_bf16_stg.py @@ -572,6 +572,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -601,6 +605,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -624,7 +633,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_normal_bf16_stg.py index 72706f27417..0baecf8f94d 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_normal_bf16_stg.py @@ -572,6 +572,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -601,6 +605,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -624,7 +633,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_swap_bf16_stg.py index f15dbaa2127..d77b4f0246c 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8/sources/sm120_fc2_non_gated_swap_bf16_stg.py @@ -571,6 +571,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -600,6 +604,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -623,7 +632,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/cudnn_frost_selected_kernels.json b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/cudnn_frost_selected_kernels.json index 1ed79925f2a..ce06e301888 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/cudnn_frost_selected_kernels.json +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/cudnn_frost_selected_kernels.json @@ -7,6 +7,14 @@ "swiglu_step", "situ" ], + "gated_param": [ + "swiglu", + "geglu", + "geglu_tanh", + "swiglu_step", + "situ", + "swiglu_param" + ], "non_gated": [ "gelu", "identity", @@ -8492,7 +8500,7 @@ ] }, { - "activation_group": "gated", + "activation_group": "gated_param", "arch": "sm_120a", "stage": "fc1", "variants": [ @@ -8526,9 +8534,9 @@ "56112e941b571de6", "7a4225debec6eddc", "41dc5fce5f0554b1", - "93767dfc83987731", + "41fd5323feb58adf", "29578b1cccc93958", - "d0718f1e440ae3aa", + "6f88d1729ff4d53e", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8546,6 +8554,18 @@ "e3b0c44298fc1c14", "4328fc11f5cb9a36" ], + "swiglu_param": [ + "2754a26085d00492", + "3b18f23975a95a83", + "6c203c815404dd02", + "6ed3d2a160090752", + "b6f78000c3b8d982", + "f15d042ef49953d5", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "52d76c29c152ab54" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -8574,13 +8594,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8597,7 +8617,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_gated_normal_bf16_stg.py", - "sha256": "ad6a58f2068901cb60e1741627094c43e0f4f9e428c9110501a1cda45ad274fd" + "sha256": "c014b694785a06741a3a05cc53553b87612f81fa702305e4cd27b55b34af6896" }, "store_policy": "stg", "tile_parameters": { @@ -8779,9 +8799,9 @@ "56112e941b571de6", "7a4225debec6eddc", "41dc5fce5f0554b1", - "425e3214fd5ae955", + "3eed40bd0e026d1e", "29578b1cccc93958", - "ffca1f68b5a79844", + "7a376b4eefc70b3c", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8799,6 +8819,18 @@ "e3b0c44298fc1c14", "b5d0f3bef19393b9" ], + "swiglu_param": [ + "2754a26085d00492", + "3b18f23975a95a83", + "6c203c815404dd02", + "0494fa69c080922a", + "b6f78000c3b8d982", + "066608727c187eb0", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "567045b85e6e674e" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -8828,13 +8860,13 @@ "intermediate_dtype": "bfloat16", "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "float8_e4m3fn", "output_scale_layout": "segmented_F8_128x4", @@ -8852,11 +8884,13 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_gated_normal_quantized_stg.py", - "sha256": "8393aa0eff0b74b33ca815ca823329a71d1eb55d80f65363deabbbe925b24282" + "sha256": "1900f462708738ce1463a2e41692220c6d039a1adc5a8d4eb9dea69b1500268d" }, "store_policy": "stg", "tile_parameters": { "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps2x4": { + "cgrp_tile_mnk": "(128, 128, 128)", + "cta_tile_mnk": "(128, 128, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "4", @@ -8882,8 +8916,6 @@ "sfb_tma_box_mn": "1" }, "CONFIG_sm120_32x256x128_16x16x32_cluster1x1_warps1x8": { - "cgrp_tile_mnk": "(32, 256, 128)", - "cta_tile_mnk": "(32, 256, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "num_b_operands": "2", @@ -8893,8 +8925,6 @@ "sA_tma_bytes": "4096" }, "CONFIG_sm120_32x256x128_16x16x32_cluster1x1_warps2x4": { - "cgrp_tile_mnk": "(32, 256, 128)", - "cta_tile_mnk": "(32, 256, 128)", "gemm_a_idx": "(0, 0)", "gemm_b_idx": "(0, 1)", "mma_size_m": "1", @@ -8954,10 +8984,10 @@ "56112e941b571de6", "eb00231ff480d1e5", "41dc5fce5f0554b1", - "93767dfc83987731", + "41fd5323feb58adf", "7bea654a536d6bf5", "29578b1cccc93958", - "d0718f1e440ae3aa", + "6f88d1729ff4d53e", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8976,6 +9006,19 @@ "e3b0c44298fc1c14", "a2ba14ff465ff35e" ], + "swiglu_param": [ + "2754a26085d00492", + "d5f025c6d2c92c97", + "6c203c815404dd02", + "6ed3d2a160090752", + "dc101805d6303043", + "b6f78000c3b8d982", + "f15d042ef49953d5", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "6f5ddd4860594594" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -9005,13 +9048,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -9028,7 +9071,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_gated_swap_bf16_stg.py", - "sha256": "501d7f1e8e023fc963ab5cbc20ab3f8bc7fbe9c9cba4121b6ec5d082fd3d26af" + "sha256": "ec44c63520928119437d02ee662eb88fd3774ada1a3378779cadcae21d4a6648" }, "store_policy": "stg", "tile_parameters": { @@ -9492,13 +9535,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -9515,7 +9558,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_non_gated_normal_bf16_stg.py", - "sha256": "fb7ed918e296b9a97e5b1f24bf8187372f0ba9d7d0175905c46253f28cfeacbb" + "sha256": "2bfe5c522fb3bcd30bba3d3334d1f407f35874c055f71eace499fbc4424ebbc3" }, "store_policy": "stg", "tile_parameters": { @@ -9670,13 +9713,13 @@ "intermediate_dtype": "bfloat16", "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "float8_e4m3fn", "output_scale_layout": "segmented_F8_128x4", @@ -9694,11 +9737,13 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_non_gated_normal_quantized_stg.py", - "sha256": "95d06c2e341c60431d799e04c78fea7ce6445a8cdb6e47912df0440c146d49bf" + "sha256": "d4d9ef0215dc83af858f1cd2a40dbff0a963476e0092e22d03705c71bb48a84c" }, "store_policy": "stg", "tile_parameters": { "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps2x4": { + "cgrp_tile_mnk": "(128, 128, 128)", + "cta_tile_mnk": "(128, 128, 128)", "mma_size_m": "4", "sB_packed_elems": "8192", "sB_tma_bytes": "8192", @@ -9706,6 +9751,8 @@ "sfb_tma_box_mn": "1" }, "CONFIG_sm120_128x128x128_16x16x32_cluster1x1_warps4x2": { + "cgrp_tile_mnk": "(128, 128, 128)", + "cta_tile_mnk": "(128, 128, 128)", "mma_size_n": "4", "sB_packed_elems": "8192", "sB_tma_bytes": "8192", @@ -9729,8 +9776,6 @@ "mma_size_n": "8" }, "CONFIG_sm120_32x256x128_16x16x32_cluster1x1_warps1x8": { - "cgrp_tile_mnk": "(32, 256, 128)", - "cta_tile_mnk": "(32, 256, 128)", "sA_packed_elems": "4096", "sA_tma_bytes": "4096" }, @@ -9798,13 +9843,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -9821,7 +9866,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_non_gated_swap_bf16_stg.py", - "sha256": "1089202c813f48fba32279c60690dd047ce056d9d24328e312cd2983ce3706b8" + "sha256": "1deb85b231a99b57f7577b6c8072a73a5c0d3039c34b09bbe37ec6b5d0873ba9" }, "store_policy": "stg", "tile_parameters": { @@ -10164,13 +10209,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -10187,7 +10232,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc2_non_gated_normal_bf16_stg.py", - "sha256": "b490619653bcc5b1ce091defedc284714ad4031ee597e86a4ec53c9c1774acac" + "sha256": "e36235a8827b4f86d6c257542d18ed7f3746c4ec64c596dbbc3baf7cf35e7724" }, "store_policy": "stg", "tile_parameters": { @@ -10305,13 +10350,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -10328,7 +10373,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc2_non_gated_swap_bf16_stg.py", - "sha256": "ca1254f0f7dca69e0eabd98badb311df951bddcb70023b68cfb539a9bde7063b" + "sha256": "799a25e845e62f32c50810856cbe29535c3e86e50c490b3f5a93539404cb466b" }, "store_policy": "stg", "tile_parameters": { @@ -10623,7 +10668,9 @@ "00c5ee6c27d970e0": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "015dae03e6eec8cf": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "047483e0c9f06583": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", + "0494fa69c080922a": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _q0_src = (_r_10).to(cutlass.Float32)\n", "0560352581e4245f": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", + "066608727c187eb0": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _q0_src = (_r_10).to(cutlass.Float32)\n", "097c496041ff0970": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _tap_0 = (_r_3).to(cutlass.BFloat16)\n", "09f974d1874b91c5": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", "0b80a44ad5fe8c55": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", @@ -10635,6 +10682,7 @@ "1dfc658c5c45d2e7": " gate_scale,\n linear_scale,\n", "1ffaeb97f374dcf1": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "22ca0d8f6cd1c1b0": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", + "2754a26085d00492": " swiglu_limit: cute.Tensor,\n swiglu_alpha: cute.Tensor,\n swiglu_lower: cute.Tensor,\n swiglu_beta: cute.Tensor,\n", "29578b1cccc93958": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "2c22943ab1bd98cf": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "2f39a5d044e06acf": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", @@ -10643,8 +10691,10 @@ "330cb022858093b7": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", "33aa3fa477b45f2f": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n vec_out = (_r_11).to(cutlass.BFloat16)\n", "378b021ec68df0b8": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", + "3b18f23975a95a83": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", + "3eed40bd0e026d1e": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "41dc5fce5f0554b1": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", - "425e3214fd5ae955": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", + "41fd5323feb58adf": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "430a595659be8727": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", "4328fc11f5cb9a36": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "458018ab4b9c29cd": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", @@ -10652,7 +10702,9 @@ "4babc9e5e7349459": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "50ed9889665429ca": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", "5265cd619c815f29": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "52d76c29c152ab54": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "56112e941b571de6": " gate_scale: cute.Tensor,\n linear_scale: cute.Tensor,\n", + "567045b85e6e674e": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "59195a7258d63f15": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _q0_src = (_r_3).to(cutlass.Float32)\n", "5a54255f3f367e7e": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "5a8b7730d3a750c0": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", @@ -10660,9 +10712,14 @@ "61f8f91878afa3b0": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "663a4532415b8373": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "66ed7d44845792bf": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", + "6c203c815404dd02": " _op_2 = cute.math.min(_c_2_a, (cutlass.full_like(_c_2_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", + "6ed3d2a160090752": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _tap_0 = (_r_10).to(cutlass.BFloat16)\n", + "6f5ddd4860594594": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "6f88d1729ff4d53e": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "7023effb53be67be": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "7928eb53a5feaeef": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", "79bdcfbd725e3cdd": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", + "7a376b4eefc70b3c": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "7a4225debec6eddc": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", "7b83af455c81d304": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "7bea654a536d6bf5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", @@ -10681,7 +10738,6 @@ "90b5af501e75575f": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "90c60ce1df6c64fb": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n vec_out = (_r_3).to(cutlass.BFloat16)\n", "91ac55be27ccb00f": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", - "93767dfc83987731": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "9399417aa8da2382": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _q0_src = (_r_1).to(cutlass.Float32)\n", "941640545da1616a": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n vec_out = (_r_4).to(cutlass.BFloat16)\n", "949c2cdeadcfe453": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _tap_0 = (_r_4).to(cutlass.BFloat16)\n", @@ -10700,6 +10756,8 @@ "b47ef4c40992b214": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "b5d0f3bef19393b9": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "b6b4da9f62018b7e": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", + "b6f78000c3b8d982": " _op_2 = cute.math.min(_c_2_a, (cutlass.full_like(_c_2_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", + "b932e17e4b585343": " swiglu_limit,\n swiglu_alpha,\n swiglu_lower,\n swiglu_beta,\n", "bafb9857135fcddb": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "bcb449ff12f30cc5": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _q0_src = (_r_4).to(cutlass.Float32)\n", "bda8354eac44ce9a": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", @@ -10721,15 +10779,18 @@ "d41b4a8d60116de9": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "d43d926df1c8e8ed": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "d446e6dcb27a27f0": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", + "d5f025c6d2c92c97": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "d6aa0440273ef725": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", "d74c1b876c2737f8": " fake_gate_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_linear_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "d8077193e3e6cac7": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "d8b53fe287d3344e": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", + "dc101805d6303043": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "e3159eb3d0eec96f": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "e3b0c44298fc1c14": "", "e48523ab7fbb239b": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n vec_out = (_r_11).to(cutlass.BFloat16)\n", "e493bd216c43ecc2": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "e4ab4f7e50c8c2ee": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "e75461c9caab2535": " fake_swiglu_limit = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_alpha = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_lower = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_beta = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "e84fea998f6de219": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "ea8c14561c6061cf": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "eb00231ff480d1e5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", @@ -10739,11 +10800,11 @@ "ef1a34ae80069ad6": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "f0a06d037b34e7e2": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "f0cd44c3d5366545": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", + "f15d042ef49953d5": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _tap_0 = (_r_10).to(cutlass.BFloat16)\n", "f174b87147f89859": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _q0_src = (_r_2).to(cutlass.Float32)\n", "f71332286643d5dc": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "fc9346af90b1bc7e": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "ff9d405c8dbd6211": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _q0_src = (_r_6).to(cutlass.Float32)\n", - "ffca1f68b5a79844": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _q0_src = (_r_11).to(cutlass.Float32)\n", "ffde0b8549b9703a": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n" }, "parameter_bases": { @@ -11823,9 +11884,9 @@ "b_tma_swizzle": "_tma.TensorMapSwizzle.s64b", "block_size": "32", "cd_dtype": "cutlass.Float8E4M3FN", - "cgrp_tile_mnk": "(128, 128, 128)", + "cgrp_tile_mnk": "(32, 256, 128)", "cluster_shape_mnk": "(1, 1, 1)", - "cta_tile_mnk": "(128, 128, 128)", + "cta_tile_mnk": "(32, 256, 128)", "epi_chunk_elems": "32", "external_counter_reset": "False", "fake_dequant_a": "False", @@ -11949,7 +12010,8 @@ "producer": { "name": "cudnn_frost", "revisions": [ - "0d72d59efd21401c7e13d606584fef2544e45934" + "0d72d59efd21401c7e13d606584fef2544e45934", + "886d7176dbe724a05f0c3955c0a87abb326afd59" ] }, "schema_version": 4, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_bf16_stg.py index 286a44bcc84..e0cc313c498 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_bf16_stg.py @@ -577,6 +577,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -606,6 +610,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -629,7 +638,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_quantized_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_quantized_stg.py index f9452faf423..f506a35dd11 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_quantized_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_normal_quantized_stg.py @@ -615,6 +615,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -644,6 +648,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -667,7 +676,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_swap_bf16_stg.py index ad473b5db72..ea330e431ac 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_gated_swap_bf16_stg.py @@ -576,6 +576,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -605,6 +609,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -628,7 +637,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py index 6a23b137690..07199d2691b 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py @@ -573,6 +573,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -602,6 +606,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -625,7 +634,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_quantized_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_quantized_stg.py index 5cca44e768a..1a34b9fd932 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_quantized_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_normal_quantized_stg.py @@ -611,6 +611,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -640,6 +644,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -663,7 +672,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py index d3efbd69e44..5bf392a672a 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py @@ -572,6 +572,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -601,6 +605,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -624,7 +633,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py index 72706f27417..0baecf8f94d 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py @@ -572,6 +572,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -601,6 +605,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -624,7 +633,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py index f15dbaa2127..d77b4f0246c 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/mxfp8_mxfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py @@ -571,6 +571,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -600,6 +604,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -623,7 +632,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/cudnn_frost_selected_kernels.json b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/cudnn_frost_selected_kernels.json index 6818eb4fc51..499fab5dda9 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/cudnn_frost_selected_kernels.json +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/cudnn_frost_selected_kernels.json @@ -7,6 +7,14 @@ "swiglu_step", "situ" ], + "gated_param": [ + "swiglu", + "geglu", + "geglu_tanh", + "swiglu_step", + "situ", + "swiglu_param" + ], "non_gated": [ "gelu", "identity", @@ -8075,7 +8083,7 @@ ] }, { - "activation_group": "gated", + "activation_group": "gated_param", "arch": "sm_120a", "stage": "fc1", "variants": [ @@ -8109,9 +8117,9 @@ "56112e941b571de6", "7a4225debec6eddc", "41dc5fce5f0554b1", - "93767dfc83987731", + "41fd5323feb58adf", "29578b1cccc93958", - "d0718f1e440ae3aa", + "6f88d1729ff4d53e", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8129,6 +8137,18 @@ "e3b0c44298fc1c14", "4328fc11f5cb9a36" ], + "swiglu_param": [ + "2754a26085d00492", + "3b18f23975a95a83", + "6c203c815404dd02", + "6ed3d2a160090752", + "b6f78000c3b8d982", + "f15d042ef49953d5", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "52d76c29c152ab54" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -8158,13 +8178,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8181,7 +8201,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_gated_normal_bf16_stg.py", - "sha256": "ad6a58f2068901cb60e1741627094c43e0f4f9e428c9110501a1cda45ad274fd" + "sha256": "c014b694785a06741a3a05cc53553b87612f81fa702305e4cd27b55b34af6896" }, "store_policy": "stg", "tile_parameters": { @@ -8347,10 +8367,10 @@ "56112e941b571de6", "eb00231ff480d1e5", "41dc5fce5f0554b1", - "93767dfc83987731", + "41fd5323feb58adf", "7bea654a536d6bf5", "29578b1cccc93958", - "d0718f1e440ae3aa", + "6f88d1729ff4d53e", "56112e941b571de6", "1dfc658c5c45d2e7", "d74c1b876c2737f8", @@ -8369,6 +8389,19 @@ "e3b0c44298fc1c14", "a2ba14ff465ff35e" ], + "swiglu_param": [ + "2754a26085d00492", + "d5f025c6d2c92c97", + "6c203c815404dd02", + "6ed3d2a160090752", + "dc101805d6303043", + "b6f78000c3b8d982", + "f15d042ef49953d5", + "2754a26085d00492", + "b932e17e4b585343", + "e75461c9caab2535", + "6f5ddd4860594594" + ], "swiglu_step": [ "e3b0c44298fc1c14", "e3b0c44298fc1c14", @@ -8399,13 +8432,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8422,7 +8455,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_gated_swap_bf16_stg.py", - "sha256": "501d7f1e8e023fc963ab5cbc20ab3f8bc7fbe9c9cba4121b6ec5d082fd3d26af" + "sha256": "ec44c63520928119437d02ee662eb88fd3774ada1a3378779cadcae21d4a6648" }, "store_policy": "stg", "tile_parameters": { @@ -8552,13 +8585,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8575,7 +8608,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc1_non_gated_normal_bf16_stg.py", - "sha256": "fb7ed918e296b9a97e5b1f24bf8187372f0ba9d7d0175905c46253f28cfeacbb" + "sha256": "2bfe5c522fb3bcd30bba3d3334d1f407f35874c055f71eace499fbc4424ebbc3" }, "store_policy": "stg", "tile_parameters": { @@ -8710,13 +8743,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8733,7 +8766,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc1_non_gated_swap_bf16_stg.py", - "sha256": "1089202c813f48fba32279c60690dd047ce056d9d24328e312cd2983ce3706b8" + "sha256": "1deb85b231a99b57f7577b6c8072a73a5c0d3039c34b09bbe37ec6b5d0873ba9" }, "store_policy": "stg", "tile_parameters": { @@ -8842,13 +8875,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8865,7 +8898,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd.py", "source": { "path": "sources/sm120_fc2_non_gated_normal_bf16_stg.py", - "sha256": "b490619653bcc5b1ce091defedc284714ad4031ee597e86a4ec53c9c1774acac" + "sha256": "e36235a8827b4f86d6c257542d18ed7f3746c4ec64c596dbbc3baf7cf35e7724" }, "store_policy": "stg", "tile_parameters": { @@ -8952,13 +8985,13 @@ }, "k": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "n": { "max": 1048576, - "min": 128, - "multiple_of": 128 + "min": 32, + "multiple_of": 32 }, "output_dtype": "bfloat16", "s": { @@ -8975,7 +9008,7 @@ "producer_template": "sm120_moe_grouped_block_scale_matmul_fwd_swap_ab.py", "source": { "path": "sources/sm120_fc2_non_gated_swap_bf16_stg.py", - "sha256": "ca1254f0f7dca69e0eabd98badb311df951bddcb70023b68cfb539a9bde7063b" + "sha256": "799a25e845e62f32c50810856cbe29535c3e86e50c490b3f5a93539404cb466b" }, "store_policy": "stg", "tile_parameters": { @@ -9031,6 +9064,7 @@ "21d179ea1255c905": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_quant_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "22ca0d8f6cd1c1b0": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "23fa3103e243d044": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_scale, fake_quant_scale, stream=_fake_stream, options=frost_compile_options)\n", + "2754a26085d00492": " swiglu_limit: cute.Tensor,\n swiglu_alpha: cute.Tensor,\n swiglu_lower: cute.Tensor,\n swiglu_beta: cute.Tensor,\n", "29578b1cccc93958": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "2c22943ab1bd98cf": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", "2ec2696f684a0c65": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _c_7_a = (_r_6).to(cutlass.Float32)\n _op_7 = _c_7_a * (cutlass.full_like(_c_7_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_7).to(cutlass.Float32)\n", @@ -9038,7 +9072,9 @@ "32bc744e8ef2c644": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "330cb022858093b7": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", "33aa3fa477b45f2f": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n vec_out = (_r_11).to(cutlass.BFloat16)\n", + "3b18f23975a95a83": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "41dc5fce5f0554b1": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", + "41fd5323feb58adf": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "430a595659be8727": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", "4328fc11f5cb9a36": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "48aae67ad287824c": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _c_5_a = (_r_6).to(cutlass.Float32)\n _op_5 = _c_5_a * (cutlass.full_like(_c_5_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_5).to(cutlass.Float32)\n _q0_frg = cute.TensorSSA(_q0_src.ir_value(), (16, 2), cutlass.Float32)\n", @@ -9047,6 +9083,7 @@ "4babc9e5e7349459": " _op_2 = _c_2_a / (cutlass.full_like(_c_2_a, _aux_gate_scale_pre.to(cutlass.Float32)))\n", "4bc86e8635ead774": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _c_3_a = (_r_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_3).to(cutlass.Float32)\n", "5265cd619c815f29": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "52d76c29c152ab54": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "53c31371f53a4193": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _c_5_a = (_r_4).to(cutlass.Float32)\n _op_5 = _c_5_a * (cutlass.full_like(_c_5_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_5).to(cutlass.Float32)\n _q0_frg = cute.TensorSSA(_q0_src.ir_value(), (16, 4), cutlass.Float32)\n", "56112e941b571de6": " gate_scale: cute.Tensor,\n linear_scale: cute.Tensor,\n", "5a54255f3f367e7e": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", @@ -9057,7 +9094,11 @@ "66ed7d44845792bf": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "67627a7118f2f293": " _r3 = _c_3_a\n _rk3 = cute.math.max(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True) + cutlass.full_like(_r3, cutlass.Float32(1.0)) * cute.math.min(_r3, cutlass.full_like(_r3, cutlass.Float32(0.0)), propagate_nan=True)\n _ru3 = cute.math.min(_rk3, cutlass.full_like(_r3, cutlass.Float32(7.0)), propagate_nan=True)\n _op_3 = _ru3\n _c_4_a = (_op_1).to(cutlass.Float32)\n _r4 = _c_4_a\n _rl4 = cute.math.max(_r4, cutlass.full_like(_r4, cutlass.Float32(-7.0)), propagate_nan=True)\n _ru4 = cute.math.min(_rl4, cutlass.full_like(_r4, cutlass.Float32(7.0)), propagate_nan=True)\n _op_4 = _ru4\n _c_5_a = (_op_3).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * (cutlass.full_like(_c_6_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_6 = (_op_6).to(cutlass.BFloat16)\n _c_7_a = (_r_6).to(cutlass.Float32)\n _op_7 = _c_7_a * (cutlass.full_like(_c_7_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_7).to(cutlass.Float32)\n _q0_frg = cute.TensorSSA(_q0_src.ir_value(), (16, 4), cutlass.Float32)\n", "692e10f29b457932": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_quant_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", + "6c203c815404dd02": " _op_2 = cute.math.min(_c_2_a, (cutlass.full_like(_c_2_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", "6d7f7819e1bce68e": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _c_12_a = (_r_11).to(cutlass.Float32)\n _op_12 = _c_12_a * (cutlass.full_like(_c_12_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_12).to(cutlass.Float32)\n", + "6ed3d2a160090752": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _tap_0 = (_r_10).to(cutlass.BFloat16)\n", + "6f5ddd4860594594": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_swiglu_limit, fake_swiglu_alpha, fake_swiglu_lower, fake_swiglu_beta, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", + "6f88d1729ff4d53e": " _th3 = _c_3_a\n _thq3 = _th3 * _th3\n _op_3 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th3) < cutlass.full_like(_th3, cutlass.Float32(0.125)), _th3 + _th3 * _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.3333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(0.13333333333333333)) + _thq3 * (cutlass.full_like(_thq3, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th3, cutlass.Float32(1.0)) - cutlass.full_like(_th3, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th3 * cutlass.full_like(_th3, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th3, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _th8 = _c_8_a\n _thq8 = _th8 * _th8\n _op_8 = cutlass._mlir_helpers.vector.where(cute.math.abs(_th8) < cutlass.full_like(_th8, cutlass.Float32(0.125)), _th8 + _th8 * _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.3333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(0.13333333333333333)) + _thq8 * (cutlass.full_like(_thq8, cutlass.Float32(-0.05396825396825397))))), (cutlass.full_like(_th8, cutlass.Float32(1.0)) - cutlass.full_like(_th8, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_th8 * cutlass.full_like(_th8, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_th8, cutlass.Float32(1.0)), approx=True, ftz=True)))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "7a4225debec6eddc": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", "7b83af455c81d304": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "7bea654a536d6bf5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", @@ -9074,7 +9115,6 @@ "90b5af501e75575f": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "90c60ce1df6c64fb": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n vec_out = (_r_3).to(cutlass.BFloat16)\n", "90ec8b8b14b0f52a": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _c_2_a = (_r_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_2).to(cutlass.Float32)\n", - "93767dfc83987731": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _tap_0 = (_r_11).to(cutlass.BFloat16)\n", "941640545da1616a": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n vec_out = (_r_4).to(cutlass.BFloat16)\n", "949c2cdeadcfe453": " _c_3_b = (_op_1).to(cutlass.Float32)\n _op_3 = _c_3_a * _c_3_b\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = _c_4_a * (cutlass.full_like(_c_4_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_4 = (_op_4).to(cutlass.BFloat16)\n _tap_0 = (_r_4).to(cutlass.BFloat16)\n", "97f8d0d45d764855": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", @@ -9092,6 +9132,8 @@ "b20cc3d369b742d6": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _tap_0 = (_r_1).to(cutlass.BFloat16)\n", "b47ef4c40992b214": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "b6b4da9f62018b7e": " _e_x2 = _c_2_a\n _e_erf2 = cute.math.erf(_e_x2 * cutlass.full_like(_e_x2, cutlass.Float32(0.7071067811865475)))\n _op_2 = cutlass.full_like(_e_x2, cutlass.Float32(0.5)) * _e_x2 * (cutlass.full_like(_e_x2, cutlass.Float32(1.0)) + _e_erf2)\n", + "b6f78000c3b8d982": " _op_2 = cute.math.min(_c_2_a, (cutlass.full_like(_c_2_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n", + "b932e17e4b585343": " swiglu_limit,\n swiglu_alpha,\n swiglu_lower,\n swiglu_beta,\n", "bafb9857135fcddb": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", "bba8f95f6f1a1f45": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n _c_2_a = (_r_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_2).to(cutlass.Float32)\n", "bd5c7dc233c7f33c": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_c_tap_1, fake_gate_alpha, fake_up_alpha, fake_scale, fake_quant_scale, stream=_fake_stream, options=frost_compile_options)\n", @@ -9113,9 +9155,11 @@ "d4121e9106cf71c9": " _op_1 = cute.math.max(_c_1_a, cutlass.full_like(_c_1_a, cutlass.Float32(0.0)), propagate_nan=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _c_2_b = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * _c_2_b\n _c_3_a = (_op_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_3 = (_op_3).to(cutlass.BFloat16)\n _tap_0 = (_r_3).to(cutlass.BFloat16)\n", "d41b4a8d60116de9": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "d43d926df1c8e8ed": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", + "d5f025c6d2c92c97": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "d74c1b876c2737f8": " fake_gate_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_linear_scale = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "d8077193e3e6cac7": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "d9c0aa76e4dc09b7": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _c_12_a = (_r_11).to(cutlass.Float32)\n _op_12 = _c_12_a * (cutlass.full_like(_c_12_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_12).to(cutlass.Float32)\n _q0_frg = cute.TensorSSA(_q0_src.ir_value(), (16, 4), cutlass.Float32)\n", + "dc101805d6303043": " _aux_swiglu_limit_ptr = swiglu_limit.iterator.raw_ptr()\n _aux_swiglu_limit_pre = (_aux_swiglu_limit_ptr + 0).load()\n _aux_swiglu_alpha_ptr = swiglu_alpha.iterator.raw_ptr()\n _aux_swiglu_alpha_pre = (_aux_swiglu_alpha_ptr + 0).load()\n _aux_swiglu_lower_ptr = swiglu_lower.iterator.raw_ptr()\n _aux_swiglu_lower_pre = (_aux_swiglu_lower_ptr + 0).load()\n _aux_swiglu_beta_ptr = swiglu_beta.iterator.raw_ptr()\n _aux_swiglu_beta_pre = (_aux_swiglu_beta_ptr + 0).load()\n", "e03f004a82142c3f": " _q0_amax2 = _q0_abs[None, 2].reduce(cute.ReductionOp.MAX, cutlass.Float32(0.0), 0)\n _q0_sf2 = _q0_amax2 * _q0_rl\n _q0_scale2 = (_q0_sf2).to(cutlass.Float8E4M3FN)\n _q0_up2 = (_q0_scale2).to(cutlass.Float32)\n _q0_inv2 = cute.math.min(cute.arch.rcp_approx(_q0_up2), cutlass.Float32(3.402823466e38))\n _q0_out[32], _q0_out[33] = cute.arch.mul_packed_f32x2((_q0_src[32], _q0_src[33]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[34], _q0_out[35] = cute.arch.mul_packed_f32x2((_q0_src[34], _q0_src[35]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[36], _q0_out[37] = cute.arch.mul_packed_f32x2((_q0_src[36], _q0_src[37]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[38], _q0_out[39] = cute.arch.mul_packed_f32x2((_q0_src[38], _q0_src[39]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[40], _q0_out[41] = cute.arch.mul_packed_f32x2((_q0_src[40], _q0_src[41]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[42], _q0_out[43] = cute.arch.mul_packed_f32x2((_q0_src[42], _q0_src[43]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[44], _q0_out[45] = cute.arch.mul_packed_f32x2((_q0_src[44], _q0_src[45]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_out[46], _q0_out[47] = cute.arch.mul_packed_f32x2((_q0_src[46], _q0_src[47]), (_q0_inv2, _q0_inv2), rnd=\"rn\", ftz=False)\n _q0_amax3 = _q0_abs[None, 3].reduce(cute.ReductionOp.MAX, cutlass.Float32(0.0), 0)\n _q0_sf3 = _q0_amax3 * _q0_rl\n _q0_scale3 = (_q0_sf3).to(cutlass.Float8E4M3FN)\n _q0_up3 = (_q0_scale3).to(cutlass.Float32)\n _q0_inv3 = cute.math.min(cute.arch.rcp_approx(_q0_up3), cutlass.Float32(3.402823466e38))\n _q0_out[48], _q0_out[49] = cute.arch.mul_packed_f32x2((_q0_src[48], _q0_src[49]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[50], _q0_out[51] = cute.arch.mul_packed_f32x2((_q0_src[50], _q0_src[51]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[52], _q0_out[53] = cute.arch.mul_packed_f32x2((_q0_src[52], _q0_src[53]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[54], _q0_out[55] = cute.arch.mul_packed_f32x2((_q0_src[54], _q0_src[55]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[56], _q0_out[57] = cute.arch.mul_packed_f32x2((_q0_src[56], _q0_src[57]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[58], _q0_out[59] = cute.arch.mul_packed_f32x2((_q0_src[58], _q0_src[59]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[60], _q0_out[61] = cute.arch.mul_packed_f32x2((_q0_src[60], _q0_src[61]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n _q0_out[62], _q0_out[63] = cute.arch.mul_packed_f32x2((_q0_src[62], _q0_src[63]), (_q0_inv3, _q0_inv3), rnd=\"rn\", ftz=False)\n", "e0de65eb3a85b8cc": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _c_3_a = (_r_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_3).to(cutlass.Float32)\n", "e3159eb3d0eec96f": " _op_2 = _c_2_a * cute.math.rcp(cutlass.full_like(_c_2_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_2_a * cutlass.full_like(_c_2_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n", @@ -9123,6 +9167,7 @@ "e48523ab7fbb239b": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n vec_out = (_r_11).to(cutlass.BFloat16)\n", "e493bd216c43ecc2": " _op_1 = _c_1_a * (cutlass.full_like(_c_1_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_1 = (_op_1).to(cutlass.BFloat16)\n vec_out = (_r_1).to(cutlass.BFloat16)\n", "e4ab4f7e50c8c2ee": " _op_1 = _c_1_a * cute.math.rcp(cutlass.full_like(_c_1_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_1_a * cutlass.full_like(_c_1_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _tap_0 = (_r_2).to(cutlass.BFloat16)\n", + "e75461c9caab2535": " fake_swiglu_limit = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_alpha = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_lower = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n fake_swiglu_beta = cute.runtime.make_fake_tensor(cutlass.Float32, (1, 1, 1), stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), assumed_align=4)\n", "ea8c14561c6061cf": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_gate_alpha, fake_up_alpha, fake_gate_scale, fake_linear_scale, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "eb00231ff480d1e5": " _aux_gate_scale_ptr = gate_scale.iterator.raw_ptr()\n _aux_gate_scale_pre = (_aux_gate_scale_ptr + 0).load()\n _aux_linear_scale_ptr = linear_scale.iterator.raw_ptr()\n _aux_linear_scale_pre = (_aux_linear_scale_ptr + 0).load()\n", "eb356af7fe12f3f2": " _op_3 = (cutlass.full_like(_c_3_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_3_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_3_a * cutlass.full_like(_c_3_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_3_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = (cutlass.full_like(_c_4_a, _aux_gate_scale_pre.to(cutlass.Float32))) * _c_4_a\n _c_5_a = (_op_0).to(cutlass.Float32)\n _op_5 = cute.math.rcp(cutlass.full_like(_c_5_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_5_a * cutlass.full_like(_c_5_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_6_a = (_op_4).to(cutlass.Float32)\n _c_6_b = (_op_5).to(cutlass.Float32)\n _op_6 = _c_6_a * _c_6_b\n _c_7_a = (_op_1).to(cutlass.Float32)\n _op_7 = _c_7_a / (cutlass.full_like(_c_7_a, _aux_linear_scale_pre.to(cutlass.Float32)))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = (cutlass.full_like(_c_8_a, cutlass.Float32(1.0)) - cutlass.full_like(_c_8_a, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_c_8_a * cutlass.full_like(_c_8_a, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_c_8_a, cutlass.Float32(1.0)), approx=True, ftz=True))\n _c_9_a = (_op_8).to(cutlass.Float32)\n _op_9 = (cutlass.full_like(_c_9_a, _aux_linear_scale_pre.to(cutlass.Float32))) * _c_9_a\n _c_10_a = (_op_6).to(cutlass.Float32)\n _c_10_b = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * _c_10_b\n _c_11_a = (_op_10).to(cutlass.Float32)\n _op_11 = _c_11_a * (cutlass.full_like(_c_11_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_11 = (_op_11).to(cutlass.BFloat16)\n _c_5_a = (_r_11).to(cutlass.Float32)\n _op_5 = _c_5_a * (cutlass.full_like(_c_5_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_5).to(cutlass.Float32)\n _q0_frg = cute.TensorSSA(_q0_src.ir_value(), (16, 2), cutlass.Float32)\n", @@ -9131,6 +9176,7 @@ "ef1a34ae80069ad6": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_gate_alpha, fake_up_alpha, fake_scale, fake_c_0, stream=_fake_stream, options=frost_compile_options)\n", "f0a06d037b34e7e2": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_a_tma_workspace, fake_token_index, fake_a_0, fake_b_0, fake_b_1, fake_sfa_0, fake_sfb_0, fake_sfb_1, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "f0cd44c3d5366545": " _g_x2 = _c_2_a\n _g_inner2 = cutlass.full_like(_g_x2, cutlass.Float32(0.7978845608028654)) * (_g_x2 + cutlass.full_like(_g_x2, cutlass.Float32(0.044715)) * _g_x2 * _g_x2 * _g_x2)\n _g_tanh2 = (cutlass.full_like(_g_inner2, cutlass.Float32(1.0)) - cutlass.full_like(_g_inner2, cutlass.Float32(2.0)) * cute.math.rcp(cute.math.exp2(_g_inner2 * cutlass.full_like(_g_inner2, cutlass.Float32(2.8853900817779268)), fastmath=True) + cutlass.full_like(_g_inner2, cutlass.Float32(1.0)), approx=True, ftz=True))\n _op_2 = cutlass.full_like(_g_x2, cutlass.Float32(0.5)) * _g_x2 * (cutlass.full_like(_g_x2, cutlass.Float32(1.0)) + _g_tanh2)\n", + "f15d042ef49953d5": " _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_swiglu_alpha_pre.to(cutlass.Float32)))\n _c_4_a = (_op_3).to(cutlass.Float32)\n _op_4 = cute.math.rcp(cutlass.full_like(_c_4_a, cutlass.Float32(1.0)) + cute.math.exp2(-_c_4_a * cutlass.full_like(_c_4_a, cutlass.Float32(1.4426950408889634)), fastmath=True), approx=True, ftz=True)\n _c_5_a = (_op_2).to(cutlass.Float32)\n _c_5_b = (_op_4).to(cutlass.Float32)\n _op_5 = _c_5_a * _c_5_b\n _c_6_a = (_op_1).to(cutlass.Float32)\n _op_6 = cute.math.min(_c_6_a, (cutlass.full_like(_c_6_a, _aux_swiglu_limit_pre.to(cutlass.Float32))))\n _c_7_a = (_op_6).to(cutlass.Float32)\n _op_7 = cute.math.max(_c_7_a, (cutlass.full_like(_c_7_a, _aux_swiglu_lower_pre.to(cutlass.Float32))))\n _c_8_a = (_op_7).to(cutlass.Float32)\n _op_8 = _c_8_a + (cutlass.full_like(_c_8_a, _aux_swiglu_beta_pre.to(cutlass.Float32)))\n _c_9_a = (_op_5).to(cutlass.Float32)\n _c_9_b = (_op_8).to(cutlass.Float32)\n _op_9 = _c_9_a * _c_9_b\n _c_10_a = (_op_9).to(cutlass.Float32)\n _op_10 = _c_10_a * (cutlass.full_like(_c_10_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_10 = (_op_10).to(cutlass.BFloat16)\n _tap_0 = (_r_10).to(cutlass.BFloat16)\n", "f71332286643d5dc": " return cute.compile(_host, problem_size, fake_first_token_offset, fake_tma_workspace, fake_token_index, fake_a_0, fake_a_1, fake_b_0, fake_sfa_0, fake_sfa_1, fake_sfb_0, fake_c_tap_0, fake_gate_alpha, fake_up_alpha, fake_scale, stream=_fake_stream, options=frost_compile_options)\n", "fc9346af90b1bc7e": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n vec_out = (_r_2).to(cutlass.BFloat16)\n", "fe9e7fbce8367c35": " _e_x1 = _c_1_a\n _e_erf1 = cute.math.erf(_e_x1 * cutlass.full_like(_e_x1, cutlass.Float32(0.7071067811865475)))\n _op_1 = cutlass.full_like(_e_x1, cutlass.Float32(0.5)) * _e_x1 * (cutlass.full_like(_e_x1, cutlass.Float32(1.0)) + _e_erf1)\n _c_2_a = (_op_1).to(cutlass.Float32)\n _op_2 = _c_2_a * (cutlass.full_like(_c_2_a, _aux_scale_pre.to(cutlass.Float32)))\n _r_2 = (_op_2).to(cutlass.BFloat16)\n _c_3_a = (_r_2).to(cutlass.Float32)\n _op_3 = _c_3_a * (cutlass.full_like(_c_3_a, _aux_quant_scale_pre.to(cutlass.Float32)))\n _q0_src = (_op_3).to(cutlass.Float32)\n", @@ -10207,7 +10253,8 @@ "producer": { "name": "cudnn_frost", "revisions": [ - "0d72d59efd21401c7e13d606584fef2544e45934" + "0d72d59efd21401c7e13d606584fef2544e45934", + "886d7176dbe724a05f0c3955c0a87abb326afd59" ] }, "schema_version": 4, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_normal_bf16_stg.py index 286a44bcc84..e0cc313c498 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_normal_bf16_stg.py @@ -577,6 +577,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -606,6 +610,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -629,7 +638,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_swap_bf16_stg.py index ad473b5db72..ea330e431ac 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_gated_swap_bf16_stg.py @@ -576,6 +576,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -605,6 +609,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -628,7 +637,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py index 6a23b137690..07199d2691b 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_normal_bf16_stg.py @@ -573,6 +573,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -602,6 +606,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -625,7 +634,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py index d3efbd69e44..5bf392a672a 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc1_non_gated_swap_bf16_stg.py @@ -572,6 +572,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -601,6 +605,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -624,7 +633,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py index 72706f27417..0baecf8f94d 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_normal_bf16_stg.py @@ -572,6 +572,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_m = cutlass.Int32(0) + # The 128-row SF blocks the host sizes the segmented SFA and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(M), num_groups) + sf_capacity = sf_active + (cutlass.Int32(M) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -601,6 +605,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the M routed + # rows, so a bad table drops rows instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(M)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[0]) * tiles_along_n my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -624,7 +633,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_m + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_m + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py index f15dbaa2127..d77b4f0246c 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/artifacts/nvfp4/sources/sm120_fc2_non_gated_swap_bf16_stg.py @@ -571,6 +571,10 @@ def frost_template_kernel( start_linear_idx = cutlass.Int32(0) total_tiles = cutlass.Int32(0) start_sf_block_n = cutlass.Int32(0) + # The 128-token SF blocks the host sizes the segmented token SF and grouped row-quant scale blobs for + # (fusion_ir.segmented_row_scale_capacity_rows). + sf_active = cutlass.min(cutlass.Int32(N), num_groups) + sf_capacity = sf_active + (cutlass.Int32(N) - sf_active) // 128 scan_base = cutlass.Int32(0) group_idx = cutlass.Int32(0) group_begin = cutlass.Int32(0) @@ -600,6 +604,11 @@ def frost_template_kernel( if my_group != 0: my_begin = cutlass.Int32(first_token_arr[my_group]) my_end = cutlass.Int32(first_token_arr[my_group + 1]) + # The offsets are device data the host never checks: clamp the group into the N routed + # tokens, so a bad table drops tokens instead of storing past the output or reading past + # token_index / token_ks, and its tile and SF block counts cannot go negative. + my_end = cutlass.min(cutlass.max(my_end, cutlass.Int32(0)), cutlass.Int32(N)) + my_begin = cutlass.min(cutlass.max(my_begin, cutlass.Int32(0)), my_end) my_tiles = cute.ceil_div(my_end - my_begin, cgrp_tile_mnk[1]) * tiles_along_m my_sf_blocks = cute.ceil_div(my_end - my_begin, 128) prefix_tiles = my_tiles @@ -623,7 +632,9 @@ def frost_template_kernel( prefix_tiles += prefix_delta prefix_sf += prefix_sf_delta my_start = start_linear_idx + prefix_tiles - my_tiles - my_sf_start = start_sf_block_n + prefix_sf - my_sf_blocks + # The clamped groups of a decreasing table can still overlap and so push the SF prefix past the + # capacity: keep every SF segment inside it. Groups of a valid table always fit. + my_sf_start = cutlass.min(start_sf_block_n + prefix_sf - my_sf_blocks, sf_capacity - my_sf_blocks) thread_succeed = nvvm.vote_sync( full_warp_mask, linear_idx < my_start + my_tiles, diff --git a/flashinfer/fused_moe/backends/cudnn_frost/bf16/moe.py b/flashinfer/fused_moe/backends/cudnn_frost/bf16/moe.py index 2c6bb5948da..31f59d024b3 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/bf16/moe.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/bf16/moe.py @@ -31,7 +31,12 @@ seal_kernels, ) from ..capabilities import require_compiler -from ..activations import ACTIVATIONS, activation_name +from ..activations import ( + ACTIVATIONS, + DEFAULT_LAUNCH_SCALARS, + activation_name, + launch_scalars, +) from ..support import SUPPORTED_ARCHS, weight_view from ..compiler import require_moe_assembler from ..shortlist import capture_profiles @@ -43,6 +48,18 @@ _WEIGHT_KEYS = frozenset(("fc1_expert_weights", "fc2_expert_weights")) _TAG = "cudnn_frost-bf16-moe-v2" +# Native executor slots: output, the output scale, SiTU's two scales and +# swiglu_param's alpha, beta, limit and -limit (the executor's scalar buffer). +_TAIL_SLOTS = { + "output": -1, + "scale": 0, + "gate_scale": 1, + "linear_scale": 2, + "swiglu_alpha": 3, + "swiglu_beta": 4, + "swiglu_limit": 5, + "swiglu_lower": 6, +} _UNSUPPORTED = "cuDNN Frost BF16 MoE kernels require SM107a or SM120a" @@ -188,6 +205,7 @@ def __init__( executor_cache, fma_tactic=None, validate_fma_sources=None, + scalars=DEFAULT_LAUNCH_SCALARS, ): first = tuple(runtime.bind_kernel(k, device) for k in first) second = tuple(runtime.bind_kernel(k, device) for k in second) @@ -198,7 +216,7 @@ def __init__( module = _module(runtime._arch_for(device)) for a, b in product(first, second): key = (_TAG, a.tactic, b.tactic) - executor_key = (hidden, intermediate, experts, topk, False, key) + executor_key = (hidden, intermediate, experts, topk, False, scalars, key) plan = executor_cache.get(executor_key) if plan is None: factory = ( @@ -215,15 +233,11 @@ def __init__( a.workspace_bytes, b.workspace_bytes, a.gated, - [ - {"output": -1, "scale": 0, "gate_scale": 1, "linear_scale": 2}[ - name - ] - for name in a.launch_tail - ], + [_TAIL_SLOTS[name] for name in a.launch_tail], a.swap_ab, b.swap_ab, False, + *scalars, ) executor_cache[executor_key] = plan # Packed states retain Modules whose run Functions borrow them. @@ -231,7 +245,15 @@ def __init__( self.launches[key] = plan["run"] required = max(required, plan["workspace_size"](tokens)) if fma_tactic is not None: - executor_key = (hidden, intermediate, experts, topk, False, fma_tactic) + executor_key = ( + hidden, + intermediate, + experts, + topk, + False, + scalars, + fma_tactic, + ) plan = executor_cache.get(executor_key) if plan is None: if validate_fma_sources is not None: @@ -260,6 +282,7 @@ def __init__( False, False, True, + *scalars, ) if validate_fma_sources is not None: validate_fma_sources() @@ -326,6 +349,15 @@ def _check_support(self): } ) ) + if not any( + kernel.arch == arch and kernel.activation == name + for kernel in runtime._discover(roots) + ): + # E.g. swiglu_param, whose scalar-parameterized sources are packaged + # for SM120 only. + raise NotImplementedError( + f"cuDNN Frost BF16 MoE has no {name} kernels for {arch}" + ) require_graph_resource_retention() if arch == "sm_107a": require_moe_assembler() @@ -358,6 +390,7 @@ def _build(self): executor_cache, metadata.fma_tactic, metadata.validate_fma_sources, + launch_scalars(activation), ), ) self._plans = self._engine.plans diff --git a/flashinfer/fused_moe/backends/cudnn_frost/bf16/runtime.py b/flashinfer/fused_moe/backends/cudnn_frost/bf16/runtime.py index b089289589c..9b5ef657a78 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/bf16/runtime.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/bf16/runtime.py @@ -15,7 +15,12 @@ import torch from .....autotuner import OptimizationProfile, TunableRunner -from ..activations import ACTIVATIONS, contract_activation, is_gated +from ..activations import ( + KNOWN_ACTIVATIONS, + SWIGLU_PARAM_SCALARS, + contract_activation, + is_gated, +) from ..runtime import ( _MANIFEST as _MANIFEST, @@ -113,7 +118,7 @@ def _read_root( if not op.startswith("grouped_gemm1_"): continue activation = contract_activation(raw.get("contract", {})) - if activation not in ACTIVATIONS or op != f"grouped_gemm1_{activation}": + if activation not in KNOWN_ACTIVATIONS or op != f"grouped_gemm1_{activation}": raise RuntimeError(f"Invalid FC1 activation contract: {artifact_id}") if any( raw["contract"].get(key) != "bfloat16" @@ -137,11 +142,11 @@ def _read_root( "multiple of 128 bytes" ) launch_tail = tuple(raw.get("launch", {}).get("tail", ())) - expected_tail = ( - {"output", "scale", "gate_scale", "linear_scale"} - if activation == "situ" - else {"output", "scale"} - ) + expected_tail = {"output", "scale"} + if activation == "situ": + expected_tail |= {"gate_scale", "linear_scale"} + elif activation == "swiglu_param": + expected_tail |= set(SWIGLU_PARAM_SCALARS) if set(launch_tail) != expected_tail or len(launch_tail) != len(expected_tail): raise RuntimeError( f"{path}: kernel {artifact_id!r} has unsupported launch tail " diff --git a/flashinfer/fused_moe/backends/cudnn_frost/bf16/support.py b/flashinfer/fused_moe/backends/cudnn_frost/bf16/support.py index df7a93b3b12..f055a274dfd 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/bf16/support.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/bf16/support.py @@ -42,7 +42,10 @@ def large_bf16_moe(config, act, arch): activation_name(config.activation) except NotImplementedError: return False - return moe_geometry(config, act) + # SM120 BF16 has no scale-factor layout; its kernels take whole + # 32-element blocks (catalog contracts) and routing moves 8-element + # packs, so e.g. gpt-oss's H = I = 2880 is admitted there. + return moe_geometry(config, act, multiple=32 if arch == 120 else 128) return False diff --git a/flashinfer/fused_moe/backends/cudnn_frost/catalog.py b/flashinfer/fused_moe/backends/cudnn_frost/catalog.py index 33b8849e18b..5adbd619a7f 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/catalog.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/catalog.py @@ -38,8 +38,12 @@ def generic_contract(raw: dict[str, Any]) -> dict[str, Any]: """Conservative common bounds for the frozen symbolic grouped-MoE recipes. 128-element N/K alignment covers BF16 vector/TMA stores and every supported - block-scale layout. Offsets remain unaligned, int32 and one per routed group; - aligned-offset and split-K recipes require a separate launch contract. + block-scale layout of the SM107 recipes. The SM120 warp-MMA recipes zero-fill + K tails through TMA and bound-check N at their STG stores, so they need only + whole 32-element blocks: 16-byte E2M1 TMA rows, one MX scale block and one + fused-quantization block (e.g. gpt-oss's 2880 = 90 x 32). Offsets remain + unaligned, int32 and one per routed group; aligned-offset and split-K recipes + require a separate launch contract. """ parameters = raw.get("source", {}).get("parameters", {}) if parameters.get("version") != 1: @@ -53,10 +57,11 @@ def generic_contract(raw: dict[str, Any]) -> dict[str, Any]: ): raise ValueError("unsupported generic Frost grouped-MoE recipe") contract = deepcopy(raw["contract"]) + multiple = 32 if raw.get("arch") == "sm_120a" else 128 contract.update( s={"min": 1}, - n={"min": 128, "max": 1 << 20, "multiple_of": 128}, - k={"min": 128, "max": 1 << 20, "multiple_of": 128}, + n={"min": multiple, "max": 1 << 20, "multiple_of": multiple}, + k={"min": multiple, "max": 1 << 20, "multiple_of": multiple}, experts={"min": 1, "max": 1024}, groups={"min": 1, "max": 1024}, ) diff --git a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_bf16.cu b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_bf16.cu index 146a8dbd659..63fd961daf5 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_bf16.cu +++ b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_bf16.cu @@ -22,6 +22,22 @@ namespace { size_t align128(size_t n) { return (n + 127) / 128 * 128; } void checked(cudaError_t err) { TVM_FFI_ICHECK_EQ(err, cudaSuccess) << cudaGetErrorString(err); } +// Scalar FC1 inputs, refreshed by the routing launch: the output scale (1), SiTU's +// gate/linear scales (4, 25) and swiglu_param's alpha, beta, limit and -limit. +struct ActivationScalars { + float alpha, beta, limit, lower; +}; +constexpr int kScalarSlots = 7; +__device__ __forceinline__ void write_scalars(float* scale, ActivationScalars act) { + scale[0] = 1.f; + scale[1] = 4.f; + scale[2] = 25.f; + scale[3] = act.alpha; + scale[4] = act.beta; + scale[5] = act.limit; + scale[6] = act.lower; +} + __global__ void histogram(const int32_t* ids, int32_t* counts, int rows, int experts) { for (int64_t r = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; r < rows; r += int64_t(blockDim.x) * gridDim.x) { @@ -33,16 +49,14 @@ __global__ void histogram(const int32_t* ids, int32_t* counts, int rows, int exp } __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors, int experts, - float* scale) { + float* scale, ActivationScalars act) { int start = 0; for (int e = 0; e < experts; ++e) { offsets[e] = cursors[e] = start; start += counts[e]; } offsets[experts] = start; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } __global__ void gather(const __nv_bfloat16* x, const int32_t* ids, int32_t* cursors, @@ -69,8 +83,9 @@ __global__ void gather(const __nv_bfloat16* x, const int32_t* ids, int32_t* curs // Each CTA computes a stable destination, independently of CTA scheduling. template __global__ void route_small(const __nv_bfloat16* x, const int32_t* ids, int32_t* offsets, - int32_t* mapping, __nv_bfloat16* grouped, float* scale, int rows, - int hidden, int topk, int experts, int32_t* token_index = nullptr) { + int32_t* mapping, __nv_bfloat16* grouped, float* scale, + ActivationScalars act, int rows, int hidden, int topk, int experts, + int32_t* token_index = nullptr) { const int r = blockIdx.x; const int lane = threadIdx.x % 32; const int warp = threadIdx.x / 32; @@ -103,9 +118,7 @@ __global__ void route_small(const __nv_bfloat16* x, const int32_t* ids, int32_t* if (r < experts) offsets[r] = partial[1][0] + partial[1][1] + partial[1][2] + partial[1][3]; if (r == 0) { offsets[experts] = rows; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } } __syncthreads(); @@ -168,7 +181,7 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { CudnnFrostMoeRunner(Array fc1, Array fc2, int64_t hidden, int64_t intermediate, int64_t experts, int64_t topk, int device, size_t scratch1, size_t scratch2, bool gated, Array tail, bool swap1, - bool swap2, bool fma, bool gather) + bool swap2, bool fma, bool gather, ActivationScalars act) : fc1_(std::move(fc1)), fc2_(std::move(fc2)), h_(hidden), @@ -183,7 +196,8 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { swap1_(swap1), swap2_(swap2), fma_(fma), - gather_(gather) {} + gather_(gather), + act_(act) {} const char* kind() const final { return "cudnn_frost_bf16_moe_runner"; } Optional GetFunction(const tvm::ffi::String& name) final { @@ -230,7 +244,7 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { layout.offsets = reserve(fma_ ? 0 : (e_ + 1) * 4); layout.cursors = reserve(fma_ ? 0 : e_ * 4); layout.mapping = reserve(fma_ ? 0 : s * 4); - layout.scale = reserve(fma_ ? 0 : 3 * sizeof(float)); + layout.scale = reserve(fma_ ? 0 : kScalarSlots * sizeof(float)); // Routing metadata must not move with a stage candidate's scratch size. layout.token_index = reserve(gather_ ? s * sizeof(int32_t) : 0); layout.scratch = reserve(std::max(scratch1_, scratch2_)); @@ -305,16 +319,16 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { if (gather_ && s <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), expert_ids, offsets, mapping, gx, - scale, s, h_, k_, e_, token_index); + scale, act_, s, h_, k_, e_, token_index); } else if (s <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), expert_ids, offsets, mapping, gx, - scale, s, h_, k_, e_); + scale, act_, s, h_, k_, e_); } else { checked(cudaMemsetAsync(counts, 0, e_ * 4, stream)); histogram<<((s + 255) / 256, 1024), 256, 0, stream>>>(expert_ids, counts, s, e_); - prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, e_, scale); + prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, e_, scale, act_); if (gather_) { frost_gather::indices<<((s + 255) / 256, 1024), 256, 0, stream>>>( expert_ids, cursors, mapping, token_index, nullptr, s, k_, e_); @@ -347,8 +361,8 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { DLTensor down{w2.data_ptr(), device_, 3, dl_bfloat16, w2shape, w2stride, 0}; DLTensor first{offsets, device_, 1, dl_int32, eshape, unit, 0}; DLTensor desc{scratch, device_, 1, dl_int64, dshape, unit, 0}; - DLTensor factors[3]; - for (int j = 0; j < 3; ++j) + DLTensor factors[kScalarSlots]; + for (int j = 0; j < kScalarSlots; ++j) factors[j] = DLTensor{scale + j, device_, 3, dl_float32, scale_shape, scale_stride, 0}; int64_t mshape_sw[]{i_, s, 1}, mstride_sw[]{1, i_, s * i_}; int64_t yshape_sw[]{h_, s, 1}, ystride_sw[]{1, h_, s * h_}; @@ -359,16 +373,21 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { // Descriptor storage is populated by the kernel before it is consumed. // AnyView borrows the TensorView descriptor; keep every descriptor alive // until CallPacked returns (temporary TensorViews would dangle). - std::array tensors{TensorView(&first), - TensorView(&desc), - TensorView(swap1_ ? &gate : &tx), - TensorView(swap1_ ? (gated_ ? &up : &tx) : &gate), - TensorView(swap1_ ? &tx : &up), - TensorView(fc1_out), - TensorView(&factors[0]), - TensorView(&factors[1]), - TensorView(&factors[2])}; - tvm::ffi::AnyView args[12]; + std::array tensors{ + TensorView(&first), + TensorView(&desc), + TensorView(swap1_ ? &gate : &tx), + TensorView(swap1_ ? (gated_ ? &up : &tx) : &gate), + TensorView(swap1_ ? &tx : &up), + TensorView(fc1_out), + TensorView(&factors[0]), + TensorView(&factors[1]), + TensorView(&factors[2]), + TensorView(&factors[3]), + TensorView(&factors[4]), + TensorView(&factors[5]), + TensorView(&factors[6])}; + tvm::ffi::AnyView args[16]; int argc = 0; args[argc++] = problem1; for (int j = 0; j < 2; ++j) args[argc++] = tensors[j]; @@ -409,13 +428,15 @@ class CudnnFrostMoeRunner final : public tvm::ffi::ModuleObj { bool gated_; Array tail_; bool swap1_, swap2_, fma_, gather_; + ActivationScalars act_; }; template Module make_runner(Array fc1, Array fc2, int64_t hidden, int64_t intermediate, int64_t experts, int64_t topk, int64_t device, int64_t scratch1, int64_t scratch2, bool gated, Array tail, bool swap1, bool swap2, - bool fma) { + bool fma, double swiglu_alpha, double swiglu_beta, double swiglu_limit, + double swiglu_lower) { TVM_FFI_ICHECK(!GatherInput || !fma); // Limit the private ABI to practical int32-indexed dimensions and avoid // overflow in the workspace/launch descriptors even for malformed callers. @@ -434,18 +455,21 @@ Module make_runner(Array fc1, Array fc2, int64_t hidden, int } else { TVM_FFI_ICHECK_EQ(fc1.size(), 1); TVM_FFI_ICHECK_EQ(fc2.size(), 1); - TVM_FFI_ICHECK(tail.size() == 2 || tail.size() == 4); + TVM_FFI_ICHECK(tail.size() == 2 || tail.size() == 4 || tail.size() == 6); int mask = 0; for (auto slot : tail) { - TVM_FFI_ICHECK(slot >= -1 && slot <= 2); + TVM_FFI_ICHECK(slot >= -1 && slot < kScalarSlots); TVM_FFI_ICHECK_EQ(mask & (1 << (slot + 1)), 0); mask |= 1 << (slot + 1); } - TVM_FFI_ICHECK(mask == 3 || mask == 15); + // output + scale, plus SiTU's two scales or swiglu_param's four scalars. + TVM_FFI_ICHECK(mask == 3 || mask == 15 || (gated && mask == (3 | 240))); } return Module(tvm::ffi::make_object( std::move(fc1), std::move(fc2), hidden, intermediate, experts, topk, device, scratch1, - scratch2, gated, std::move(tail), swap1, swap2, fma, GatherInput)); + scratch2, gated, std::move(tail), swap1, swap2, fma, GatherInput, + ActivationScalars{float(swiglu_alpha), float(swiglu_beta), float(swiglu_limit), + float(swiglu_lower)})); } } // namespace diff --git a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8.cu b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8.cu index 02fda8e1c33..a7076c1fa59 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8.cu +++ b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8.cu @@ -23,6 +23,22 @@ constexpr DLDataType dl_e8m0{kDLFloat8_e8m0fnu, 8, 1}; size_t align128(size_t n) { return (n + 127) / 128 * 128; } void checked(cudaError_t err) { TVM_FFI_ICHECK_EQ(err, cudaSuccess) << cudaGetErrorString(err); } +// Scalar FC1 inputs, refreshed by the routing launch: the output scale (1), SiTU's +// gate/linear scales (4, 25) and swiglu_param's alpha, beta, limit and -limit. +struct ActivationScalars { + float alpha, beta, limit, lower; +}; +constexpr int kScalarSlots = 7; +__device__ __forceinline__ void write_scalars(float* scale, ActivationScalars act) { + scale[0] = 1.f; + scale[1] = 4.f; + scale[2] = 25.f; + scale[3] = act.alpha; + scale[4] = act.beta; + scale[5] = act.limit; + scale[6] = act.lower; +} + __global__ void histogram(const int32_t* ids, int32_t* counts, int rows, int experts) { for (int64_t r = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; r < rows; r += int64_t(blockDim.x) * gridDim.x) { @@ -85,7 +101,7 @@ __device__ int64_t sf_index(int row, int col, int columns) { } __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors, - int32_t* sf_offsets, int experts, float* scale) { + int32_t* sf_offsets, int experts, float* scale, ActivationScalars act) { int start = 0, sf_start = 0; for (int e = 0; e < experts; ++e) { offsets[e] = cursors[e] = start; @@ -94,9 +110,7 @@ __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors sf_start += (counts[e] + 127) / 128 * 128; } offsets[experts] = start; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } __global__ void gather(const uint8_t* x, const uint8_t* input_sf, const int32_t* ids, @@ -134,8 +148,8 @@ template __global__ void route_small(const uint8_t* x, const uint8_t* input_sf, const int32_t* ids, int32_t* offsets, int32_t* sf_offsets, int32_t* mapping, int32_t* row_experts, uint8_t* grouped, uint8_t* sf, float* scale, - int rows, int hidden, int topk, int experts, bool swizzled, - int32_t* token_index = nullptr) { + ActivationScalars act, int rows, int hidden, int topk, int experts, + bool swizzled, int32_t* token_index = nullptr) { const int r = blockIdx.x; const bool active = r < rows; const int raw = active ? ids[r] : 0; @@ -178,9 +192,7 @@ __global__ void route_small(const uint8_t* x, const uint8_t* input_sf, const int } if (r == 0) { offsets[experts] = rows; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } } __syncthreads(); @@ -266,7 +278,7 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { int64_t intermediate, int64_t experts, int64_t topk, int device, size_t scratch1, size_t scratch2, bool gated, Array tail1, Array tail2, bool swap1, bool swap2, bool swizzled, bool fma, - bool gather) + bool gather, ActivationScalars act) : fc1_(std::move(fc1)), fc2_(std::move(fc2)), h_(hidden), @@ -283,7 +295,8 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { swap2_(swap2), swizzled_(swizzled), fma_(fma), - gather_(gather) { + gather_(gather), + act_(act) { quantized_ = fma_ || std::find(tail1_.begin(), tail1_.end(), 3) != tail1_.end(); } @@ -359,7 +372,7 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { layout.cursors = reserve(fma_ ? 0 : e_ * 4); layout.mapping = reserve(fma_ ? 0 : s * 4); layout.row_experts = reserve(fma_ ? 0 : s * 4); - layout.scale = reserve(fma_ ? 0 : 3 * sizeof(float)); + layout.scale = reserve(fma_ ? 0 : kScalarSlots * sizeof(float)); // Routing metadata must not move with a stage candidate's scratch size. layout.token_index = reserve(gather_ ? s * sizeof(int32_t) : 0); layout.scratch = reserve(std::max(scratch1_, scratch2_)); @@ -473,18 +486,18 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { if (gather_ && s <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), static_cast(xsf.data_ptr()), - expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, s, h_, k_, e_, - swizzled_, input_index); + expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, act_, s, h_, k_, + e_, swizzled_, input_index); } else if (s <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), static_cast(xsf.data_ptr()), - expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, s, h_, k_, e_, - swizzled_); + expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, act_, s, h_, k_, + e_, swizzled_); } else { checked(cudaMemsetAsync(counts, 0, e_ * 4, stream)); histogram<<((s + 255) / 256, 1024), 256, 0, stream>>>(expert_ids, counts, s, e_); - prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, sf_offsets, e_, scale); + prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, sf_offsets, e_, scale, act_); if (gather_) { frost_gather::indices<<((s + 255) / 256, 4096), 256, 0, stream>>>( expert_ids, cursors, mapping, input_index, row_experts, s, k_, e_); @@ -547,8 +560,8 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { DLTensor sf_out{sfm, device_, 3, dl_e8m0, sfoutshape, sfoutstride, 0}; DLTensor first{offsets, device_, 1, dl_int32, eshape, unit, 0}; DLTensor desc{scratch, device_, 1, dl_int64, dshape, unit, 0}; - DLTensor factors[3]; - for (int j = 0; j < 3; ++j) + DLTensor factors[kScalarSlots]; + for (int j = 0; j < kScalarSlots; ++j) factors[j] = DLTensor{scale + j, device_, 3, dl_float32, scalar_shape, scalar_stride, 0}; int64_t mshape_sw[]{i_, s, 1}, mstride_sw[]{1, i_, s * i_}; int64_t yshape_sw[]{h_, s, 1}, ystride_sw[]{1, h_, s * h_}; @@ -557,22 +570,27 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { // The frozen host resets its scheduler counter on every invocation. // Its kernels populate descriptor storage before consuming it. // Own the TensorView descriptors until the borrowed AnyView arguments return. - std::array tensors{TensorView(&first), - TensorView(&desc), - TensorView(swap1_ ? &gate : &tx), - TensorView(swap1_ ? (gated_ ? &up : &tx) : &gate), - TensorView(swap1_ ? &tx : &up), - TensorView(swap1_ ? &sf_gate : &sf_x), - TensorView(swap1_ ? (gated_ ? &sf_up : &sf_x) : &sf_gate), - TensorView(swap1_ ? &sf_x : &sf_up), - TensorView(quantized_ ? &tqm : (swap1_ ? &tm_sw : &tm)), - TensorView(&factors[0]), - TensorView(&factors[1]), - TensorView(&factors[2]), - TensorView(&sf_out), - TensorView(&alpha_first), - TensorView(&alpha_first)}; - tvm::ffi::AnyView args[20]; + std::array tensors{ + TensorView(&first), + TensorView(&desc), + TensorView(swap1_ ? &gate : &tx), + TensorView(swap1_ ? (gated_ ? &up : &tx) : &gate), + TensorView(swap1_ ? &tx : &up), + TensorView(swap1_ ? &sf_gate : &sf_x), + TensorView(swap1_ ? (gated_ ? &sf_up : &sf_x) : &sf_gate), + TensorView(swap1_ ? &sf_x : &sf_up), + TensorView(quantized_ ? &tqm : (swap1_ ? &tm_sw : &tm)), + TensorView(&factors[0]), + TensorView(&factors[1]), + TensorView(&factors[2]), + TensorView(&sf_out), + TensorView(&alpha_first), + TensorView(&alpha_first), + TensorView(&factors[3]), + TensorView(&factors[4]), + TensorView(&factors[5]), + TensorView(&factors[6])}; + tvm::ffi::AnyView args[24]; int argc = 0; args[argc++] = problem1; args[argc++] = tensors[0]; @@ -638,15 +656,19 @@ class CudnnFrostMxfp8MoeRunner final : public tvm::ffi::ModuleObj { bool gated_; Array tail1_, tail2_; bool swap1_, swap2_, swizzled_, quantized_, fma_, gather_; + ActivationScalars act_; }; template Module make_runner(Array fc1, Array fc2, int64_t hidden, int64_t intermediate, int64_t experts, int64_t topk, int64_t device, int64_t scratch1, int64_t scratch2, bool gated, Array tail1, Array tail2, - bool swap1, bool swap2, bool swizzled, bool fma) { + bool swap1, bool swap2, bool swizzled, bool fma, double swiglu_alpha, + double swiglu_beta, double swiglu_limit, double swiglu_lower) { TVM_FFI_ICHECK(!GatherInput || !fma); constexpr int64_t max_dim = 1 << 20; + const ActivationScalars act{float(swiglu_alpha), float(swiglu_beta), float(swiglu_limit), + float(swiglu_lower)}; TVM_FFI_ICHECK(hidden > 0 && hidden <= max_dim && intermediate > 0 && intermediate <= max_dim && experts > 0 && experts <= 1024 && topk > 0 && topk <= experts); TVM_FFI_ICHECK_EQ(hidden % 128, 0); @@ -665,21 +687,23 @@ Module make_runner(Array fc1, Array fc2, int64_t hidden, int TVM_FFI_ICHECK(scratch1 > 0 && scratch1 % 128 == 0 && scratch2 > 0 && scratch2 % 128 == 0); int mask = 0; for (auto slot : tail1) { - TVM_FFI_ICHECK(slot >= -1 && slot <= 5); + TVM_FFI_ICHECK(slot >= -1 && slot <= 9); TVM_FFI_ICHECK_EQ(mask & (1 << (slot + 1)), 0); mask |= 1 << (slot + 1); } bool quantized = mask & 16; TVM_FFI_ICHECK(!quantized || !swap1); int required = (quantized ? 19 : 3) | 32 | (gated ? 64 : 0); - TVM_FFI_ICHECK(mask == required || mask == (required | 12)); + // SiTU adds its two scales; swiglu_param (gated only) its four scalars. + TVM_FFI_ICHECK(mask == required || mask == (required | 12) || + (gated && mask == (required | 1920))); TVM_FFI_ICHECK_EQ(tail2.size(), 2); TVM_FFI_ICHECK((tail2[0] == -1 && tail2[1] == 4) || (tail2[0] == 4 && tail2[1] == -1)); } return Module(tvm::ffi::make_object( std::move(fc1), std::move(fc2), hidden, intermediate, experts, topk, device, scratch1, - scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, fma, - GatherInput)); + scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, fma, GatherInput, + act)); } } // namespace diff --git a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8_mxfp4.cu b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8_mxfp4.cu index ad974e0d7b8..07ceb3d8e25 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8_mxfp4.cu +++ b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_mxfp8_mxfp4.cu @@ -24,6 +24,22 @@ constexpr DLDataType dl_e8m0{kDLFloat8_e8m0fnu, 8, 1}; size_t align128(size_t n) { return (n + 127) / 128 * 128; } void checked(cudaError_t err) { TVM_FFI_ICHECK_EQ(err, cudaSuccess) << cudaGetErrorString(err); } +// Scalar FC1 inputs, refreshed by the routing launch: the output scale (1), SiTU's +// gate/linear scales (4, 25) and swiglu_param's alpha, beta, limit and -limit. +struct ActivationScalars { + float alpha, beta, limit, lower; +}; +constexpr int kScalarSlots = 7; +__device__ __forceinline__ void write_scalars(float* scale, ActivationScalars act) { + scale[0] = 1.f; + scale[1] = 4.f; + scale[2] = 25.f; + scale[3] = act.alpha; + scale[4] = act.beta; + scale[5] = act.limit; + scale[6] = act.lower; +} + __global__ void histogram(const int32_t* ids, int32_t* counts, int rows, int experts) { for (int64_t r = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; r < rows; r += int64_t(blockDim.x) * gridDim.x) { @@ -86,7 +102,7 @@ __device__ int64_t sf_index(int row, int col, int columns) { } __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors, - int32_t* sf_offsets, int experts, float* scale) { + int32_t* sf_offsets, int experts, float* scale, ActivationScalars act) { int start = 0, sf_start = 0; for (int e = 0; e < experts; ++e) { offsets[e] = cursors[e] = start; @@ -95,9 +111,7 @@ __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors sf_start += (counts[e] + 127) / 128 * 128; } offsets[experts] = start; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } __global__ void gather(const uint8_t* x, const uint8_t* input_sf, const int32_t* ids, @@ -135,8 +149,8 @@ template __global__ void route_small(const uint8_t* x, const uint8_t* input_sf, const int32_t* ids, int32_t* offsets, int32_t* sf_offsets, int32_t* mapping, int32_t* row_experts, uint8_t* grouped, uint8_t* sf, float* scale, - int rows, int hidden, int topk, int experts, bool swizzled, - int32_t* token_index = nullptr) { + ActivationScalars act, int rows, int hidden, int topk, int experts, + bool swizzled, int32_t* token_index = nullptr) { const int r = blockIdx.x; const bool active = r < rows; const int raw = active ? ids[r] : 0; @@ -179,9 +193,7 @@ __global__ void route_small(const uint8_t* x, const uint8_t* input_sf, const int } if (r == 0) { offsets[experts] = rows; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } } __syncthreads(); @@ -267,7 +279,7 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { int64_t intermediate, int64_t experts, int64_t topk, int device, size_t scratch1, size_t scratch2, bool gated, Array tail1, Array tail2, bool swap1, bool swap2, bool swizzled, - bool fma, bool gather) + bool fma, bool gather, ActivationScalars act) : fc1_(std::move(fc1)), fc2_(std::move(fc2)), h_(hidden), @@ -284,7 +296,8 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { swap2_(swap2), swizzled_(swizzled), fma_(fma), - gather_(gather) { + gather_(gather), + act_(act) { quantized_ = fma_ || std::find(tail1_.begin(), tail1_.end(), 3) != tail1_.end(); } @@ -360,7 +373,7 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { result.cursors = reserve(fma_ ? 0 : e_ * 4); result.mapping = reserve(fma_ ? 0 : result.rows * 4); result.row_experts = reserve(fma_ ? 0 : result.rows * 4); - result.scale = reserve(fma_ ? 0 : 3 * sizeof(float)); + result.scale = reserve(fma_ ? 0 : kScalarSlots * sizeof(float)); result.scratch = reserve(std::max(scratch1_, scratch2_)); result.token_index = reserve(gather_ ? result.rows * sizeof(int32_t) : 0); result.workspace_size = pos; @@ -477,18 +490,18 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { if (gather_ && s_ <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), static_cast(xsf.data_ptr()), - expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, s_, h_, k_, e_, - swizzled_, input_index); + expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, act_, s_, h_, k_, + e_, swizzled_, input_index); } else if (s_ <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), static_cast(xsf.data_ptr()), - expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, s_, h_, k_, e_, - swizzled_); + expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, act_, s_, h_, k_, + e_, swizzled_); } else { checked(cudaMemsetAsync(counts, 0, e_ * 4, stream)); histogram<<((s_ + 255) / 256, 1024), 256, 0, stream>>>(expert_ids, counts, s_, e_); - prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, sf_offsets, e_, scale); + prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, sf_offsets, e_, scale, act_); if (gather_) { frost_gather::indices<<((s_ + 255) / 256, 4096), 256, 0, stream>>>( expert_ids, cursors, mapping, input_index, row_experts, s_, k_, e_); @@ -549,8 +562,8 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { DLTensor sf_out{sfm, device_, 3, dl_e8m0, sfoutshape, sfoutstride, 0}; DLTensor first{offsets, device_, 1, dl_int32, eshape, unit, 0}; DLTensor desc{scratch, device_, 1, dl_int64, dshape, unit, 0}; - DLTensor factors[3]; - for (int j = 0; j < 3; ++j) + DLTensor factors[kScalarSlots]; + for (int j = 0; j < kScalarSlots; ++j) factors[j] = DLTensor{scale + j, device_, 3, dl_float32, scalar_shape, scalar_stride, 0}; int64_t mshape_sw[]{i_, s_, 1}, mstride_sw[]{1, i_, s_ * i_}; int64_t yshape_sw[]{h_, s_, 1}, ystride_sw[]{1, h_, s_ * h_}; @@ -559,22 +572,27 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { // The frozen host resets its scheduler counter on every invocation. // Its kernels populate descriptor storage before consuming it. // Own the TensorView descriptors until the borrowed AnyView arguments return. - std::array tensors{TensorView(&first), - TensorView(&desc), - TensorView(swap1_ ? &gate : &tx), - TensorView(swap1_ ? (gated_ ? &up : &tx) : &gate), - TensorView(swap1_ ? &tx : &up), - TensorView(swap1_ ? &sf_gate : &sf_x), - TensorView(swap1_ ? (gated_ ? &sf_up : &sf_x) : &sf_gate), - TensorView(swap1_ ? &sf_x : &sf_up), - TensorView(quantized_ ? &tqm : (swap1_ ? &tm_sw : &tm)), - TensorView(&factors[0]), - TensorView(&factors[1]), - TensorView(&factors[2]), - TensorView(&sf_out), - TensorView(&alpha_first), - TensorView(&alpha_first)}; - tvm::ffi::AnyView args[20]; + std::array tensors{ + TensorView(&first), + TensorView(&desc), + TensorView(swap1_ ? &gate : &tx), + TensorView(swap1_ ? (gated_ ? &up : &tx) : &gate), + TensorView(swap1_ ? &tx : &up), + TensorView(swap1_ ? &sf_gate : &sf_x), + TensorView(swap1_ ? (gated_ ? &sf_up : &sf_x) : &sf_gate), + TensorView(swap1_ ? &sf_x : &sf_up), + TensorView(quantized_ ? &tqm : (swap1_ ? &tm_sw : &tm)), + TensorView(&factors[0]), + TensorView(&factors[1]), + TensorView(&factors[2]), + TensorView(&sf_out), + TensorView(&alpha_first), + TensorView(&alpha_first), + TensorView(&factors[3]), + TensorView(&factors[4]), + TensorView(&factors[5]), + TensorView(&factors[6])}; + tvm::ffi::AnyView args[24]; int argc = 0; args[argc++] = problem1; args[argc++] = tensors[0]; @@ -640,15 +658,19 @@ class CudnnFrostMxfp8Mxfp4MoeRunner final : public tvm::ffi::ModuleObj { bool gated_; Array tail1_, tail2_; bool swap1_, swap2_, swizzled_, quantized_, fma_, gather_; + ActivationScalars act_; }; template Module make_runner(Array fc1, Array fc2, int64_t hidden, int64_t intermediate, int64_t experts, int64_t topk, int64_t device, int64_t scratch1, int64_t scratch2, bool gated, Array tail1, Array tail2, - bool swap1, bool swap2, bool swizzled, bool fma) { + bool swap1, bool swap2, bool swizzled, bool fma, double swiglu_alpha, + double swiglu_beta, double swiglu_limit, double swiglu_lower) { TVM_FFI_ICHECK(!GatherInput || !fma); constexpr int64_t max_dim = 1 << 20; + const ActivationScalars act{float(swiglu_alpha), float(swiglu_beta), float(swiglu_limit), + float(swiglu_lower)}; TVM_FFI_ICHECK(hidden > 0 && hidden <= max_dim && intermediate > 0 && intermediate <= max_dim && experts > 0 && experts <= 1024 && topk > 0 && topk <= experts); TVM_FFI_ICHECK_EQ(fc1.size(), fc2.size()); @@ -665,21 +687,23 @@ Module make_runner(Array fc1, Array fc2, int64_t hidden, int TVM_FFI_ICHECK(scratch1 > 0 && scratch1 % 128 == 0 && scratch2 > 0 && scratch2 % 128 == 0); int mask = 0; for (auto slot : tail1) { - TVM_FFI_ICHECK(slot >= -1 && slot <= 5); + TVM_FFI_ICHECK(slot >= -1 && slot <= 9); TVM_FFI_ICHECK_EQ(mask & (1 << (slot + 1)), 0); mask |= 1 << (slot + 1); } bool quantized = mask & 16; TVM_FFI_ICHECK(!quantized || !swap1); int required = (quantized ? 19 : 3) | 32 | (gated ? 64 : 0); - TVM_FFI_ICHECK(mask == required || mask == (required | 12)); + // SiTU adds its two scales; swiglu_param (gated only) its four scalars. + TVM_FFI_ICHECK(mask == required || mask == (required | 12) || + (gated && mask == (required | 1920))); TVM_FFI_ICHECK_EQ(tail2.size(), 2); TVM_FFI_ICHECK((tail2[0] == -1 && tail2[1] == 4) || (tail2[0] == 4 && tail2[1] == -1)); } return Module(tvm::ffi::make_object( std::move(fc1), std::move(fc2), hidden, intermediate, experts, topk, device, scratch1, - scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, fma, - GatherInput)); + scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, fma, GatherInput, + act)); } } // namespace diff --git a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_nvfp4.cu b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_nvfp4.cu index 13abc47b266..379318c7b80 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_nvfp4.cu +++ b/flashinfer/fused_moe/backends/cudnn_frost/csrc/moe_nvfp4.cu @@ -22,8 +22,26 @@ using tvm::ffi::Optional; namespace { constexpr DLDataType dl_fp4{kDLFloat4_e2m1fn, 4, 2}; size_t align128(size_t n) { return (n + 127) / 128 * 128; } +// 128x4 block scales pad each matrix's rows to whole 128-row atoms. +int64_t atom_rows(int64_t rows) { return (rows + 127) / 128 * 128; } void checked(cudaError_t err) { TVM_FFI_ICHECK_EQ(err, cudaSuccess) << cudaGetErrorString(err); } +// Scalar FC1 inputs, refreshed by the routing launch: the output scale (1), SiTU's +// gate/linear scales (4, 25) and swiglu_param's alpha, beta, limit and -limit. +struct ActivationScalars { + float alpha, beta, limit, lower; +}; +constexpr int kScalarSlots = 7; +__device__ __forceinline__ void write_scalars(float* scale, ActivationScalars act) { + scale[0] = 1.f; + scale[1] = 4.f; + scale[2] = 25.f; + scale[3] = act.alpha; + scale[4] = act.beta; + scale[5] = act.limit; + scale[6] = act.lower; +} + __global__ void histogram(const int32_t* ids, int32_t* counts, int rows, int experts) { for (int64_t r = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; r < rows; r += int64_t(blockDim.x) * gridDim.x) { @@ -86,7 +104,7 @@ __device__ int64_t sf_index(int row, int col, int columns) { } __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors, - int32_t* sf_offsets, int experts, float* scale) { + int32_t* sf_offsets, int experts, float* scale, ActivationScalars act) { int start = 0, sf_start = 0; for (int e = 0; e < experts; ++e) { offsets[e] = cursors[e] = start; @@ -95,9 +113,7 @@ __global__ void prefix(const int32_t* counts, int32_t* offsets, int32_t* cursors sf_start += (counts[e] + 127) / 128 * 128; } offsets[experts] = start; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } // COMBINE consumes grouped-row -> token/slot metadata. Fold the per-expert @@ -150,9 +166,10 @@ template __global__ void route_small(const uint8_t* x, const uint8_t* input_sf, const int32_t* ids, int32_t* offsets, int32_t* sf_offsets, int32_t* mapping, int32_t* row_experts, uint8_t* grouped, uint8_t* sf, float* scale, - int rows, int hidden, int topk, int experts, bool swizzled, - int32_t* slots, float* combined_scores, const float* scores, - const float* alpha, int32_t* token_index = nullptr) { + ActivationScalars act, int rows, int hidden, int topk, int experts, + bool swizzled, int32_t* slots, float* combined_scores, + const float* scores, const float* alpha, + int32_t* token_index = nullptr) { const int r = blockIdx.x; const bool active = r < rows; const int raw = active ? ids[r] : 0; @@ -196,9 +213,7 @@ __global__ void route_small(const uint8_t* x, const uint8_t* input_sf, const int } if (r == 0) { offsets[experts] = rows; - scale[0] = 1.f; - scale[1] = 4.f; - scale[2] = 25.f; + write_scalars(scale, act); } } __syncthreads(); @@ -272,7 +287,7 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { int64_t intermediate, int64_t experts, int64_t topk, int device, size_t scratch1, size_t scratch2, bool gated, Array tail1, Array tail2, bool swap1, bool swap2, bool swizzled, bool fma, - bool gather) + bool gather, ActivationScalars act) : fc1_(std::move(fc1)), fc2_(std::move(fc2)), h_(hidden), @@ -291,7 +306,8 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { swizzled_(swizzled), fma_(fma), combine_(!fma && tail2_.size() == 1), - gather_(gather) {} + gather_(gather), + act_(act) {} const char* kind() const final { return "cudnn_frost_nvfp4_moe_runner"; } Optional GetFunction(const tvm::ffi::String& name) final { @@ -369,7 +385,7 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { result.cursors = reserve(fma_ ? 0 : e_ * 4); result.mapping = reserve(fma_ ? 0 : result.rows * 4); result.row_experts = reserve(fma_ ? 0 : result.rows * 4); - result.scale = reserve(fma_ ? 0 : 3 * sizeof(float)); + result.scale = reserve(fma_ ? 0 : kScalarSlots * sizeof(float)); result.combine_slots = reserve(combine_ ? result.rows * sizeof(int32_t) : 0); result.combine_scores = reserve(combine_ ? result.rows * sizeof(float) : 0); // Routing metadata must not move with a stage candidate's scratch size. @@ -426,16 +442,20 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { tensor(scores, device_, dl_float32, {t_, k_}, 4); tensor(w1, device_, dl_uint8, {e_, (gated_ ? 2 : 1) * i_, h_ / 2}); tensor(w2, device_, dl_uint8, {e_, h_, i_ / 2}); + const int64_t i_rows = atom_rows(i_), h_rows = atom_rows(h_); const bool packed_sf1 = sf1.ndim() == 4; if (packed_sf1) - tensor(sf1, device_, dl_uint8, {(gated_ ? 2 : 1), e_, i_, h_ / 16}); + tensor(sf1, device_, dl_uint8, {(gated_ ? 2 : 1), e_, i_rows, h_ / 16}); else - tensor(sf1, device_, dl_uint8, {e_, (gated_ ? 2 : 1) * i_, h_ / 16}); + tensor(sf1, device_, dl_uint8, {e_, atom_rows((gated_ ? 2 : 1) * i_), h_ / 16}); + // Canonical [up; gate] scales split at row I only on an atom boundary. + TVM_FFI_ICHECK(packed_sf1 || !gated_ || i_ % 128 == 0) + << "gated FC1 block scales need one 128x4 matrix per branch when I % 128 != 0"; const int64_t sf1_expert_stride = sf1.stride(packed_sf1 ? 1 : 0); const int64_t sf1_gate_stride = packed_sf1 ? sf1.stride(0) : i_ * sf1.stride(1); const int64_t sf1_expert_bytes = sf1_expert_stride; const int64_t sf1_gate_bytes = gated_ ? sf1_gate_stride : 0; - tensor(sf2, device_, dl_uint8, {e_, h_, i_ / 16}); + tensor(sf2, device_, dl_uint8, {e_, h_rows, i_ / 16}); tensor(global1, device_, dl_float32, {1}, 4); tensor(global2, device_, dl_float32, {1}, 4); tensor(alpha1, device_, dl_float32, {e_}, 4); @@ -498,20 +518,22 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { if (gather_ && s_ <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), static_cast(xsf.data_ptr()), - expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, s_, h_, k_, e_, - swizzled_, combine_slots, combine_scores, static_cast(scores.data_ptr()), + expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, act_, s_, h_, k_, + e_, swizzled_, combine_slots, combine_scores, + static_cast(scores.data_ptr()), static_cast(alpha2.data_ptr()), input_index); } else if (s_ <= 512 && e_ <= 256) { route_small<<>>( static_cast(x.data_ptr()), static_cast(xsf.data_ptr()), - expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, s_, h_, k_, e_, - swizzled_, combine_slots, combine_scores, static_cast(scores.data_ptr()), + expert_ids, offsets, sf_offsets, mapping, row_experts, gx, sfx, scale, act_, s_, h_, k_, + e_, swizzled_, combine_slots, combine_scores, + static_cast(scores.data_ptr()), static_cast(alpha2.data_ptr())); } else { checked(cudaMemsetAsync(counts, 0, e_ * 4, stream)); histogram<<((s_ + 255) / 256, 1024), 256, 0, stream>>>(expert_ids, counts, s_, e_); - prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, sf_offsets, e_, scale); + prefix<<<1, 1, 0, stream>>>(counts, offsets, cursors, sf_offsets, e_, scale, act_); if (gather_) { frost_gather::indices<<((s_ + 255) / 256, 4096), 256, 0, stream>>>( expert_ids, cursors, mapping, input_index, row_experts, s_, k_, e_, combine_slots, @@ -541,8 +563,8 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { int64_t yshape[]{s_, h_, 1}, ystride[]{h_, 1, s_ * h_}; int64_t w1shape[]{i_, h_ / 2, e_}, w1stride[]{h_ / 2, 1, (gated_ ? 2 : 1) * i_ * h_ / 2}; int64_t w2shape[]{h_, i_ / 2, e_}, w2stride[]{i_ / 2, 1, h_ * i_ / 2}; - int64_t sf1shape[]{i_ * h_ / 16, 1, e_}, sf1stride[]{1, 1, sf1_expert_bytes}; - int64_t sf2shape[]{h_ * i_ / 16, 1, e_}, sf2stride[]{1, 1, h_ * i_ / 16}; + int64_t sf1shape[]{i_rows * h_ / 16, 1, e_}, sf1stride[]{1, 1, sf1_expert_bytes}; + int64_t sf2shape[]{h_rows * i_ / 16, 1, e_}, sf2stride[]{1, 1, h_rows * i_ / 16}; int64_t sfxshape[]{sf_rows_ * h_ / 16, 1, 1}, sfxstride[]{1, 1, 1}; int64_t sfmshape[]{sf_rows_ * i_ / 16, 1, 1}; int64_t eshape[]{e_ + 1}, dshape[]{int64_t(scratch1_ / 8)}, unit[]{1}; @@ -585,8 +607,8 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { scalar_shape, scalar_stride, 0}; DLTensor first{offsets, device_, 1, dl_int32, eshape, unit, 0}; DLTensor desc{scratch, device_, 1, dl_int64, dshape, unit, 0}; - DLTensor factors[3]; - for (int j = 0; j < 3; ++j) + DLTensor factors[kScalarSlots]; + for (int j = 0; j < kScalarSlots; ++j) factors[j] = DLTensor{scale + j, device_, 3, dl_float32, scalar_shape, scalar_stride, 0}; int64_t group_shape[]{e_, 1, 1}; DLTensor alpha_first{alpha1.data_ptr(), device_, 3, dl_float32, group_shape, scalar_stride, 0}; @@ -598,7 +620,7 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { // The frozen host resets its scheduler counter on every invocation, and // the kernel initializes operand/scale/output descriptors before use. // Own the TensorView descriptors until the borrowed AnyView arguments return. - std::array tensors{ + std::array tensors{ TensorView(&first), TensorView(&desc), TensorView(swap1_ ? &gate : &tx), @@ -614,8 +636,12 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { TensorView(&alpha_first), TensorView(&alpha_first), TensorView(&sf_out), - TensorView(&quant_scale)}; - tvm::ffi::AnyView args[19]; + TensorView(&quant_scale), + TensorView(&factors[3]), + TensorView(&factors[4]), + TensorView(&factors[5]), + TensorView(&factors[6])}; + tvm::ffi::AnyView args[23]; int argc = 0; args[argc++] = problem1; args[argc++] = tensors[0]; @@ -689,20 +715,26 @@ class CudnnFrostNvfp4MoeRunner final : public tvm::ffi::ModuleObj { bool gated_; Array tail1_, tail2_; bool quantized_, swap1_, swap2_, swizzled_, fma_, combine_, gather_; + ActivationScalars act_; }; template Module make_runner(Array fc1, Array fc2, int64_t hidden, int64_t intermediate, int64_t experts, int64_t topk, int64_t device, int64_t scratch1, int64_t scratch2, bool gated, Array tail1, Array tail2, - bool swap1, bool swap2, bool swizzled, bool fma) { + bool swap1, bool swap2, bool swizzled, bool fma, double swiglu_alpha, + double swiglu_beta, double swiglu_limit, double swiglu_lower) { TVM_FFI_ICHECK(!GatherInput || !fma); constexpr int64_t max_dim = 1 << 20; + const ActivationScalars act{float(swiglu_alpha), float(swiglu_beta), float(swiglu_limit), + float(swiglu_lower)}; TVM_FFI_ICHECK(hidden > 0 && hidden <= max_dim && intermediate > 0 && intermediate <= max_dim && experts > 0 && experts <= 1024 && topk > 0 && topk <= experts); TVM_FFI_ICHECK_EQ(fc1.size(), fc2.size()); - TVM_FFI_ICHECK_EQ(hidden % 128, 0); - TVM_FFI_ICHECK_EQ(intermediate % 128, 0); + // Scale rows hold whole 4-column words of 16-element blocks: H/I in 64s + // (SM120 takes e.g. gpt-oss's 2880; the SM107 runner admits 128s only). + TVM_FFI_ICHECK_EQ(hidden % 64, 0); + TVM_FFI_ICHECK_EQ(intermediate % 64, 0); if (fma) { TVM_FFI_ICHECK(!swap1 && !swap2 && scratch1 == 0 && scratch2 == 0 && tail1.empty() && tail2.empty()); @@ -711,13 +743,14 @@ Module make_runner(Array fc1, Array fc2, int64_t hidden, int (experts == 12 && hidden == 7168 && intermediate == 3072 && topk == 2 && fc1.size() == 1)); return Module(tvm::ffi::make_object( std::move(fc1), std::move(fc2), hidden, intermediate, experts, topk, device, scratch1, - scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, true, false)); + scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, true, false, + act)); } TVM_FFI_ICHECK_EQ(fc1.size(), 1); TVM_FFI_ICHECK(scratch1 > 0 && scratch1 % 128 == 0 && scratch2 > 0 && scratch2 % 128 == 0); int mask = 0; for (auto slot : tail1) { - TVM_FFI_ICHECK(slot >= 0 && slot <= 7); + TVM_FFI_ICHECK(slot >= 0 && slot <= 11); TVM_FFI_ICHECK_EQ(mask & (1 << slot), 0); mask |= 1 << slot; } @@ -726,14 +759,16 @@ Module make_runner(Array fc1, Array fc2, int64_t hidden, int required |= (1 << 6) | (1 << 7); TVM_FFI_ICHECK(!swap1); } - TVM_FFI_ICHECK(mask == required || mask == (required | 12)); + // SiTU adds its two scales; swiglu_param (gated only) its four scalars. + TVM_FFI_ICHECK(mask == required || mask == (required | 12) || + (gated && mask == (required | 3840))); TVM_FFI_ICHECK((tail2.size() == 1 && tail2[0] == 0) || (tail2.size() == 2 && ((tail2[0] == 0 && tail2[1] == 4) || (tail2[0] == 4 && tail2[1] == 0)))); return Module(tvm::ffi::make_object( std::move(fc1), std::move(fc2), hidden, intermediate, experts, topk, device, scratch1, scratch2, gated, std::move(tail1), std::move(tail2), swap1, swap2, swizzled, false, - GatherInput)); + GatherInput, act)); } } // namespace diff --git a/flashinfer/fused_moe/backends/cudnn_frost/dimension_catalog.py b/flashinfer/fused_moe/backends/cudnn_frost/dimension_catalog.py index b426b3e125b..97045fb87dd 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/dimension_catalog.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/dimension_catalog.py @@ -22,6 +22,15 @@ ) GATED = ("swiglu", "geglu", "geglu_tanh", "swiglu_step", "situ") NON_GATED = ("gelu", "identity", "relu", "relu2", "silu") +# Gated activations whose parameters are launch-time scalars. They are exported +# per architecture: a gated family that carries them expands the "gated_param" +# group (GATED plus these), other gated families keep "gated". +PARAMETERIZED_GATED = ("swiglu_param",) +GATED_PARAM = GATED + PARAMETERIZED_GATED +# Frost passes swiglu_param's scalars in first-use order, ahead of the output +# scale; write_catalog checks every exported tail against this expansion. +SWIGLU_PARAM_TAIL = ("swiglu_limit", "swiglu_alpha", "swiglu_lower", "swiglu_beta") +_GATED_GROUPS = ("gated", "gated_param") def logical_id(raw, dtype): @@ -101,13 +110,15 @@ def expand_catalog(payload: dict[str, Any]) -> dict[str, Any]: else ["gate_alpha"] + ( ["up_alpha"] - if family["activation_group"] == "gated" + if family["activation_group"] in _GATED_GROUPS else [] ) + aux ) if activation == "situ": aux = aux[:-1] + ["gate_scale", "linear_scale"] + aux[-1:] + elif activation == "swiglu_param": + aux = aux[:-1] + list(SWIGLU_PARAM_TAIL) + aux[-1:] if quantized and dtype == "nvfp4": aux.append("quant_scale") if combine: @@ -278,7 +289,9 @@ def write_catalog( if activation == "silu(gate) * up": activation = record["contract"]["activation"] = "swiglu" stage = "fc2" if record["op"].endswith("gemm2") else "fc1" - topology = "gated" if stage == "fc1" and activation in GATED else "non_gated" + topology = ( + "gated" if stage == "fc1" and activation in GATED_PARAM else "non_gated" + ) tactic = record["tactic"] tile = { k: deepcopy(v) @@ -305,6 +318,21 @@ def write_catalog( ) groups[key].append(record) + parameterized_families = { + key[:3] + for key, members in groups.items() + if any(r["contract"]["activation"] in PARAMETERIZED_GATED for r in members) + } + + def group_of(arch, stage, topology): + if topology != "gated": + return topology + return ( + "gated_param" + if (arch, stage, topology) in parameterized_families + else "gated" + ) + # Keep each producer's geometry-dependent body until its frozen fragments # have been factored into the shared template below. prepared = [] @@ -315,6 +343,7 @@ def write_catalog( groups.items(), key=lambda item: (item[0][-2] == "gather", item[0]) ): arch, stage, topology, swapped, output, store, mode, chunk = key + group = group_of(arch, stage, topology) sources: dict[str, str] = {} prefixes: set[str] = set() parameters: dict[str, dict[str, str]] = {} @@ -350,18 +379,24 @@ def write_catalog( expected = ( ("identity",) if stage == "fc2" - else GATED - if topology == "gated" - else NON_GATED + else {"gated": GATED, "gated_param": GATED_PARAM}.get(group, NON_GATED) ) if set(sources) != set(expected): if stage != "fc1": raise RuntimeError( f"missing activation source for {key}: {set(expected) - sources.keys()}" ) - # SiTU has two additional scalar arguments. Start from a common - # activation ABI when synthesizing other activation/store arms. - ordinary = sources.keys() - {"situ"} + # SiTU and the parameterized activations take additional scalar + # arguments. Start from a common activation ABI when synthesizing + # other activation/store arms; parameterized ones are never synthesized. + missing_parameterized = (set(expected) - sources.keys()) & set( + PARAMETERIZED_GATED + ) + if missing_parameterized: + raise RuntimeError( + f"missing exported {sorted(missing_parameterized)} source for {key}" + ) + ordinary = sources.keys() - {"situ", *PARAMETERIZED_GATED} if not ordinary: raise RuntimeError( "Frost activation expansion needs a common ABI source" @@ -429,7 +464,11 @@ def write_catalog( ), }, dtype=dtype, - activation_groups={"gated": list(GATED), "non_gated": list(NON_GATED)}, + activation_groups={ + "gated": list(GATED), + "non_gated": list(NON_GATED), + **({"gated_param": list(GATED_PARAM)} if parameterized_families else {}), + }, store_policies={ "stg": {"data": "stg", "scales": "stg"}, "tma": {"data": "tma", "scales": None}, @@ -463,7 +502,10 @@ def write_catalog( family_key = (arch, stage, topology) if family_key not in families: families[family_key] = dict( - arch=arch, stage=stage, activation_group=topology, variants=[] + arch=arch, + stage=stage, + activation_group=group_of(arch, stage, topology), + variants=[], ) quantized = output != "bfloat16" policy = "tmastg" if store == "tma" and quantized else store @@ -512,8 +554,19 @@ def write_catalog( ) payload["families"] = list(families.values()) expanded = expand_catalog(payload) - new_ids = {logical_id(r, dtype): r["id"] for r in expanded["kernels"]} + by_logical = {logical_id(r, dtype): r for r in expanded["kernels"]} + new_ids = {key: r["id"] for key, r in by_logical.items()} mapping = {r["id"]: new_ids[logical_id(r, dtype)] for r in records} + # The expansion recomputes every launch tail from the activation; it must + # reproduce the producer's argument order for each exported record. + for r in records: + expected_tail = r.get("launch", {}).get("tail") + actual_tail = by_logical[logical_id(r, dtype)]["launch"]["tail"] + if expected_tail is not None and list(expected_tail) != list(actual_tail): + raise RuntimeError( + f"Frost catalog launch tail {actual_tail} differs from the exported " + f"{list(expected_tail)} for {r['id']}" + ) for destination, template in pending_sources.items(): destination.parent.mkdir(parents=True, exist_ok=True) fd, temporary = tempfile.mkstemp( diff --git a/flashinfer/fused_moe/backends/cudnn_frost/export.py b/flashinfer/fused_moe/backends/cudnn_frost/export.py index 1581391e716..f6a551135ce 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/export.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/export.py @@ -13,7 +13,7 @@ from pathlib import Path from typing import Any -from .activations import ACTIVATIONS, is_gated +from .activations import KNOWN_ACTIVATIONS, SWIGLU_PARAM_SCALARS, is_gated from .manifest import MANIFEST_FILENAME, read_manifest, write_manifest from .catalog import generic_contract, kernel_id @@ -184,7 +184,7 @@ def apply_gemm_scale(tensor, name): output.set_data_type(cudnn.data_type.BFLOAT16).set_output(True) return graph activation = op.removeprefix("grouped_gemm1_") - if activation not in ACTIVATIONS: + if activation not in KNOWN_ACTIVATIONS: raise ValueError(f"unsupported export op {op!r}") up = graph.tensor( name="up_weight", @@ -256,6 +256,12 @@ def apply_gemm_scale(tensor, name): if activation == "relu2": activated = graph.mul(a=activated, b=activated, name="square") elif activation == "swiglu_step": + # min(silu(gate), 7) * clamp(up, -7, 7). cuDNN's RELU_FWD is + # x < lower ? lower + slope * (x - lower) : x, then min(., upper), with + # lower defaulting to 0: the gate keeps its negative values through an + # identity slope, and a slope-free ReLU clamps up at -7. (A slope with + # lower_clip=-7 would leak below -7; Frost before cuDNN Frontend + # e0185b3b clamped there anyway. Both packaged revisions are later.) activated = graph.relu( input=activated, negative_slope=1.0, upper_clip=7.0, name="gate_clamp" ) @@ -267,6 +273,37 @@ def apply_gemm_scale(tensor, name): upper_clip=7.0, name="up_clamp", ) + elif activation == "swiglu_param": + # SwiGLU with launch-time alpha/beta/limit (FlashInfer's SwiGLU contract): + # g = min(gate, limit), u = clamp(up, -limit, limit), + # g * sigmoid(alpha * g) * (u + beta). The lower bound arrives as -limit. + alpha, beta, limit, lower = ( + graph.tensor( + name=name, + dim=[1, 1, 1], + stride=[1, 1, 1], + data_type=cudnn.data_type.FLOAT, + ) + for name in SWIGLU_PARAM_SCALARS + ) + gate_clamped = graph.min(input0=gate_out, input1=limit, name="gate_clamp") + activated = graph.mul( + a=gate_clamped, + b=graph.sigmoid( + input=graph.mul(a=gate_clamped, b=alpha, name="gate_slope"), + name="gate_sigmoid", + ), + name="gate_swish", + ) + up_out = graph.add( + a=graph.max( + input0=graph.min(input0=up_out, input1=limit, name="up_upper"), + input1=lower, + name="up_lower", + ), + b=beta, + name="up_bias", + ) elif activation == "situ": gate_scale = graph.tensor( name="gate_scale", @@ -696,9 +733,11 @@ def _export_compiled(args, compiled, config, arch): if 0 in tma_slots else ["output", "output_scale"] + aux_names ) - expected_aux = ( - {"scale", "gate_scale", "linear_scale"} if activation == "situ" else {"scale"} - ) + expected_aux = {"scale"} + if activation == "situ": + expected_aux |= {"gate_scale", "linear_scale"} + elif activation == "swiglu_param": + expected_aux |= set(SWIGLU_PARAM_SCALARS) if dtype != "bf16": expected_aux |= {"gate_alpha"} if is_gated(activation): @@ -841,7 +880,10 @@ def main() -> None: ) parser.add_argument( "--op", - choices=(*(f"grouped_gemm1_{name}" for name in ACTIVATIONS), "grouped_gemm2"), + choices=( + *(f"grouped_gemm1_{name}" for name in KNOWN_ACTIVATIONS), + "grouped_gemm2", + ), default="grouped_gemm1_swiglu", ) parser.add_argument("--output-dir", type=Path, required=True) diff --git a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/moe.py b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/moe.py index 529ba95b9a5..ac2a9692aa0 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/moe.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/moe.py @@ -17,7 +17,12 @@ from .....fused_moe.runners import MoERunner, _validate_prerouted_inputs from .....utils import get_compute_capability from .. import runtime as common -from ..activations import ACTIVATIONS, activation_name +from ..activations import ( + ACTIVATIONS, + DEFAULT_LAUNCH_SCALARS, + activation_name, + launch_scalars, +) from ..support import SUPPORTED_ARCHS, weight_view from ..compiler import require_moe_assembler from ..cache import require_graph_resource_retention @@ -226,6 +231,7 @@ def __init__( swizzled=False, fma_tactic=None, validate_fma_sources=None, + scalars=DEFAULT_LAUNCH_SCALARS, ): self.plans, self.launches = {}, {} first = tuple(common.bind_kernel(k, device) for k in first) @@ -235,7 +241,7 @@ def __init__( module = _module(common._arch_for(device)) for a, b in product(first, second): key = (_TAG, a.tactic, b.tactic) - executor_key = (hidden, intermediate, experts, topk, swizzled, key) + executor_key = (hidden, intermediate, experts, topk, swizzled, scalars, key) plan = executor_cache.get(executor_key) if plan is None: factory = ( @@ -258,13 +264,22 @@ def __init__( b.swap_ab, swizzled, False, + *scalars, ) executor_cache[executor_key] = plan # Packed states retain Modules whose run Functions borrow them. self.plans[key], self.launches[key] = plan, plan["run"] required = max(required, plan["workspace_size"](tokens)) if fma_tactic is not None: - executor_key = (hidden, intermediate, experts, topk, swizzled, fma_tactic) + executor_key = ( + hidden, + intermediate, + experts, + topk, + swizzled, + scalars, + fma_tactic, + ) plan = executor_cache.get(executor_key) if plan is None: if validate_fma_sources is not None: @@ -302,6 +317,7 @@ def __init__( False, swizzled, True, + *scalars, ) if validate_fma_sources is not None: validate_fma_sources() @@ -426,6 +442,16 @@ def _check_support(self): } ) ) + if not any( + kernel.fc1 and kernel.arch == arch and kernel.activation == name + for root in _artifact_roots() + for kernel in runtime.discover(root) + ): + # E.g. swiglu_param, whose scalar-parameterized sources are packaged + # for SM120 only. + raise NotImplementedError( + f"cuDNN Frost MXFP8 MoE has no {name} kernels for {arch}" + ) require_graph_resource_retention() # Only the SM107 sources patch tensor maps and need the assembler floor. if arch == "sm_107a": @@ -461,6 +487,7 @@ def _build(self): swizzled, metadata.fma_tactic, metadata.validate_fma_sources, + launch_scalars(activation), ), ) self._plans = self._engine.plans diff --git a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/runtime.py b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/runtime.py index f2591bc4ba0..d9b1db59f4d 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/runtime.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8/runtime.py @@ -15,10 +15,16 @@ import torch from .. import runtime -from ..activations import ACTIVATIONS, is_gated +from ..activations import ( + ACTIVATIONS, + KNOWN_ACTIVATIONS, + SWIGLU_PARAM_SCALARS, + is_gated, + launch_scalars, +) _FC2 = "block_scale_grouped_gemm2" -_OPS = {f"block_scale_grouped_gemm1_{name}" for name in ACTIVATIONS} | {_FC2} +_OPS = {f"block_scale_grouped_gemm1_{name}" for name in KNOWN_ACTIVATIONS} | {_FC2} _artifact_cache_version = 0 _TAIL_SLOTS = { @@ -30,6 +36,10 @@ "gate_alpha": 4, "up_alpha": 5, "alpha": 4, + "swiglu_alpha": 6, + "swiglu_beta": 7, + "swiglu_limit": 8, + "swiglu_lower": 9, } @@ -146,6 +156,8 @@ def discover( expected_tail.add("up_alpha") if expected["activation"] == "situ": expected_tail |= {"gate_scale", "linear_scale"} + elif expected["activation"] == "swiglu_param": + expected_tail |= set(SWIGLU_PARAM_SCALARS) if fused: expected_tail.add("output_scale") prefix = ("output_scale",) if store == "tma" else ("output", "output_scale") @@ -363,6 +375,7 @@ def __init__( scale: torch.Tensor | None = None, gemm_scales: tuple[torch.Tensor, ...] | None = None, workspace: torch.Tensor | None = None, + activation=None, ): if tokens.device.type != "cuda": raise ValueError("MXFP8 grouped GEMM requires CUDA tensors") @@ -432,16 +445,26 @@ def __init__( raise ValueError("GEMM scales require 4B alignment") activation_scales = {} if kernel.activation == "situ": - activation = ACTIVATIONS["situ"]() + situ = ACTIVATIONS["situ"]() activation_scales = { name: torch.full( (1, 1, 1), - getattr(activation, name), + getattr(situ, name), dtype=torch.float32, device=tokens.device, ) for name in ("gate_scale", "linear_scale") } + elif kernel.activation == "swiglu_param": + # SwiGLU's launch-time alpha/beta/limit (defaults when not given). + activation_scales = { + name: torch.full( + (1, 1, 1), value, dtype=torch.float32, device=tokens.device + ) + for name, value in zip( + SWIGLU_PARAM_SCALARS, launch_scalars(activation), strict=True + ) + } if workspace is None: workspace = torch.empty( kernel.workspace_bytes, dtype=torch.uint8, device=tokens.device diff --git a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/moe.py b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/moe.py index 9da3f6b6ec6..0c49a752019 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/moe.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/moe.py @@ -17,7 +17,12 @@ from .....fused_moe.runners import MoERunner, _validate_prerouted_inputs from .....utils import get_compute_capability from .. import runtime as common -from ..activations import ACTIVATIONS, activation_name +from ..activations import ( + ACTIVATIONS, + DEFAULT_LAUNCH_SCALARS, + activation_name, + launch_scalars, +) from ..support import SUPPORTED_ARCHS, weight_view from .support import is_eligible from ..compiler import require_moe_assembler @@ -226,6 +231,7 @@ def __init__( swizzled=False, fma_tactic=None, validate_fma_sources=None, + scalars=DEFAULT_LAUNCH_SCALARS, ): self.plans, self.launches = {}, {} first = tuple(common.bind_kernel(k, device) for k in first) @@ -235,7 +241,7 @@ def __init__( module = _module(common._arch_for(device)) for a, b in product(first, second): key = (_TAG, a.tactic, b.tactic) - executor_key = (hidden, intermediate, experts, topk, swizzled, key) + executor_key = (hidden, intermediate, experts, topk, swizzled, scalars, key) plan = executor_cache.get(executor_key) if plan is None: factory = ( @@ -258,13 +264,22 @@ def __init__( b.swap_ab, swizzled, False, + *scalars, ) executor_cache[executor_key] = plan # Packed states retain Modules whose run Functions borrow them. self.plans[key], self.launches[key] = plan, plan["run"] required = max(required, plan["workspace_size"](tokens)) if fma_tactic is not None: - executor_key = (hidden, intermediate, experts, topk, swizzled, fma_tactic) + executor_key = ( + hidden, + intermediate, + experts, + topk, + swizzled, + scalars, + fma_tactic, + ) plan = executor_cache.get(executor_key) if plan is None: if validate_fma_sources is not None: @@ -302,6 +317,7 @@ def __init__( False, swizzled, True, + *scalars, ) if validate_fma_sources is not None: validate_fma_sources() @@ -427,6 +443,16 @@ def _check_support(self): } ) ) + if not any( + kernel.fc1 and kernel.arch == arch and kernel.activation == name + for root in _artifact_roots() + for kernel in runtime.discover(root) + ): + # E.g. swiglu_param, whose scalar-parameterized sources are packaged + # for SM120 only. + raise NotImplementedError( + f"cuDNN Frost MXFP8 × MXFP4 MoE has no {name} kernels for {arch}" + ) require_graph_resource_retention() # Only the SM107 sources patch tensor maps and need the assembler floor. if arch == "sm_107a": @@ -462,6 +488,7 @@ def _build(self): swizzled, metadata.fma_tactic, metadata.validate_fma_sources, + launch_scalars(activation), ), ) self._plans = self._engine.plans diff --git a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/runtime.py b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/runtime.py index 6ba0ebb6364..7c31f0da901 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/runtime.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/mxfp8_mxfp4/runtime.py @@ -16,10 +16,16 @@ import torch from .. import runtime -from ..activations import ACTIVATIONS, is_gated +from ..activations import ( + ACTIVATIONS, + KNOWN_ACTIVATIONS, + SWIGLU_PARAM_SCALARS, + is_gated, + launch_scalars, +) _FC2 = "block_scale_grouped_gemm2" -_OPS = {f"block_scale_grouped_gemm1_{name}" for name in ACTIVATIONS} | {_FC2} +_OPS = {f"block_scale_grouped_gemm1_{name}" for name in KNOWN_ACTIVATIONS} | {_FC2} _artifact_cache_version = 0 _TAIL_SLOTS = { @@ -31,6 +37,10 @@ "gate_alpha": 4, "up_alpha": 5, "alpha": 4, + "swiglu_alpha": 6, + "swiglu_beta": 7, + "swiglu_limit": 8, + "swiglu_lower": 9, } @@ -149,6 +159,8 @@ def discover( expected_tail.add("up_alpha") if expected["activation"] == "situ": expected_tail |= {"gate_scale", "linear_scale"} + elif expected["activation"] == "swiglu_param": + expected_tail |= set(SWIGLU_PARAM_SCALARS) if fused: expected_tail.add("output_scale") prefix = ("output_scale",) if store == "tma" else ("output", "output_scale") @@ -368,6 +380,7 @@ def __init__( scale: torch.Tensor | None = None, gemm_scales: tuple[torch.Tensor, ...] | None = None, workspace: torch.Tensor | None = None, + activation=None, ): if tokens.device.type != "cuda": raise ValueError("MXFP8 × MXFP4 grouped GEMM requires CUDA tensors") @@ -443,16 +456,26 @@ def __init__( raise ValueError("GEMM scales require 4B alignment") activation_scales = {} if kernel.activation == "situ": - activation = ACTIVATIONS["situ"]() + situ = ACTIVATIONS["situ"]() activation_scales = { name: torch.full( (1, 1, 1), - getattr(activation, name), + getattr(situ, name), dtype=torch.float32, device=tokens.device, ) for name in ("gate_scale", "linear_scale") } + elif kernel.activation == "swiglu_param": + # SwiGLU's launch-time alpha/beta/limit (defaults when not given). + activation_scales = { + name: torch.full( + (1, 1, 1), value, dtype=torch.float32, device=tokens.device + ) + for name, value in zip( + SWIGLU_PARAM_SCALARS, launch_scalars(activation), strict=True + ) + } if workspace is None: workspace = torch.empty( kernel.workspace_bytes, dtype=torch.uint8, device=tokens.device diff --git a/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/moe.py b/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/moe.py index 5ba20de7e8e..f32d6e68132 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/moe.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/moe.py @@ -17,7 +17,12 @@ from .....fused_moe.runners import MoERunner, _validate_prerouted_inputs from .....utils import get_compute_capability from .. import runtime as common -from ..activations import ACTIVATIONS, activation_name +from ..activations import ( + ACTIVATIONS, + DEFAULT_LAUNCH_SCALARS, + activation_name, + launch_scalars, +) from ..support import SUPPORTED_ARCHS, weight_view from .support import is_eligible from ..compiler import require_moe_assembler @@ -56,6 +61,10 @@ _UNSUPPORTED = "cuDNN Frost NVFP4 MoE kernels require SM107a or SM120a" +def _atom_rows(rows): + return (rows + 127) // 128 * 128 + + def _tensor_version(tensor): # Only scalar validation uses versions; inference scalars have no counter. # Their values must continue to satisfy the validated numerical contract. @@ -259,6 +268,7 @@ def __init__( swizzled=False, fma_tactic=None, validate_fma_sources=None, + scalars=DEFAULT_LAUNCH_SCALARS, ): self.plans, self.launches = {}, {} first = tuple(common.bind_kernel(k, device) for k in first) @@ -268,7 +278,7 @@ def __init__( module = _module(common._arch_for(device)) for a, b in product(first, second): key = (_TAG, a.tactic, b.tactic) - executor_key = (hidden, intermediate, experts, topk, swizzled, key) + executor_key = (hidden, intermediate, experts, topk, swizzled, scalars, key) plan = executor_cache.get(executor_key) if plan is None: factory = ( @@ -291,13 +301,22 @@ def __init__( b.swap_ab, swizzled, False, + *scalars, ) executor_cache[executor_key] = plan # Packed states retain Modules whose run Functions borrow them. self.plans[key], self.launches[key] = plan, plan["run"] required = max(required, plan["workspace_size"](tokens)) if fma_tactic is not None: - executor_key = (hidden, intermediate, experts, topk, swizzled, fma_tactic) + executor_key = ( + hidden, + intermediate, + experts, + topk, + swizzled, + scalars, + fma_tactic, + ) plan = executor_cache.get(executor_key) if plan is None: if validate_fma_sources is not None: @@ -335,6 +354,7 @@ def __init__( False, swizzled, True, + *scalars, ) if validate_fma_sources is not None: validate_fma_sources() @@ -467,6 +487,16 @@ def _check_support(self): } ) ) + if not any( + kernel.fc1 and kernel.arch == arch and kernel.activation == name + for root in _artifact_roots() + for kernel in runtime.discover(root) + ): + # E.g. swiglu_param, whose scalar-parameterized sources are packaged + # for SM120 only. + raise NotImplementedError( + f"cuDNN Frost NVFP4 MoE has no {name} kernels for {arch}" + ) require_graph_resource_retention() # Only the SM107 sources patch tensor maps and need the assembler floor. if arch == "sm_107a": @@ -502,6 +532,7 @@ def _build(self): swizzled, metadata.fma_tactic, metadata.validate_fma_sources, + launch_scalars(activation), ), ) self._plans = self._engine.plans @@ -547,8 +578,10 @@ def pack_inputs(self, act, weights): and 0 < i <= 1 << 20 ): raise ValueError("cuDNN Frost NVFP4 requires bounded int32 geometry") - if h % 128 or i % 128: - raise ValueError("cuDNN Frost NVFP4 requires H/I divisible by 128") + # SM120 takes H/I in 64s (4-column scale words); SM107 in 128s. + multiple = 64 if get_compute_capability(self.device) == (12, 0) else 128 + if h % multiple or i % multiple: + raise ValueError(f"cuDNN Frost NVFP4 requires H/I divisible by {multiple}") _validate_prerouted_inputs( act, t, @@ -564,10 +597,11 @@ def pack_inputs(self, act, weights): ) w1, w2, a1, sf1, alpha1, a2, sf2, alpha2 = (view[name] for name in _WEIGHT_KEYS) mult = 2 if self.config.activation.is_gated else 1 + # 128x4 block scales pad each matrix's rows to whole 128-row atoms. expected = ( (w1, torch.uint8, (e, mult * i, h // 2)), (w2, torch.uint8, (e, h, i // 2)), - (sf2, torch.uint8, (e, h, i // 16)), + (sf2, torch.uint8, (e, _atom_rows(h), i // 16)), (alpha1, torch.float32, (e,)), (alpha2, torch.float32, (e,)), ) @@ -579,13 +613,20 @@ def pack_inputs(self, act, weights): or sf1.dtype != torch.uint8 or tuple(sf1.shape) not in ( - (e, mult * i, h // 16), - (mult, e, i, h // 16), + (e, _atom_rows(mult * i), h // 16), + (mult, e, _atom_rows(i), h // 16), ) ): raise ValueError( "cuDNN Frost NVFP4 canonical weight/scale geometry mismatch" ) + if sf1.ndim == 3 and mult == 2 and i % 128: + # The gate rows start mid-atom: Frost reads up and gate separately. + raise ValueError( + "cuDNN Frost NVFP4 needs per-branch FC1 block scales when I is " + "not a multiple of 128; prepare weights with " + "CudnnFrostNvfp4Config.prepare_weights" + ) if any( value.dtype != torch.float32 or tuple(value.shape) not in ((), (1,)) for value in (a1, a2) diff --git a/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/runtime.py b/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/runtime.py index dc47cb58f89..ba707f8d4ac 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/runtime.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/runtime.py @@ -15,10 +15,16 @@ import torch from .. import runtime -from ..activations import ACTIVATIONS, is_gated +from ..activations import ( + ACTIVATIONS, + KNOWN_ACTIVATIONS, + SWIGLU_PARAM_SCALARS, + is_gated, + launch_scalars, +) _FC2 = "block_scale_grouped_gemm2" -_OPS = {f"block_scale_grouped_gemm1_{name}" for name in ACTIVATIONS} | {_FC2} +_OPS = {f"block_scale_grouped_gemm1_{name}" for name in KNOWN_ACTIVATIONS} | {_FC2} _artifact_cache_version = 0 _TAIL_SLOTS = { "output": 0, @@ -30,6 +36,10 @@ "alpha": 4, "output_scale": 6, "quant_scale": 7, + "swiglu_alpha": 8, + "swiglu_beta": 9, + "swiglu_limit": 10, + "swiglu_lower": 11, } @@ -176,6 +186,8 @@ def discover( expected_tail.add("up_alpha") if expected["activation"] == "situ": expected_tail |= {"gate_scale", "linear_scale"} + elif expected["activation"] == "swiglu_param": + expected_tail |= set(SWIGLU_PARAM_SCALARS) if fused: expected_tail |= {"output_scale", "quant_scale"} prefix = ("output_scale",) if store == "tma" else ("output", "output_scale") @@ -397,6 +409,7 @@ def __init__( scale: torch.Tensor | None = None, gemm_scales: tuple[torch.Tensor, ...] | None = None, workspace: torch.Tensor | None = None, + activation=None, ): if kernel.combines_output: raise ValueError( @@ -473,16 +486,26 @@ def __init__( ) activation_scales = {} if kernel.activation == "situ": - activation = ACTIVATIONS["situ"]() + situ = ACTIVATIONS["situ"]() activation_scales = { name: torch.full( (1, 1, 1), - getattr(activation, name), + getattr(situ, name), dtype=torch.float32, device=tokens.device, ) for name in ("gate_scale", "linear_scale") } + elif kernel.activation == "swiglu_param": + # SwiGLU's launch-time alpha/beta/limit (defaults when not given). + activation_scales = { + name: torch.full( + (1, 1, 1), value, dtype=torch.float32, device=tokens.device + ) + for name, value in zip( + SWIGLU_PARAM_SCALARS, launch_scalars(activation), strict=True + ) + } if workspace is None: workspace = torch.empty( kernel.workspace_bytes, dtype=torch.uint8, device=tokens.device diff --git a/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/support.py b/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/support.py index de741e813e9..91890e2a754 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/support.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/nvfp4/support.py @@ -23,4 +23,7 @@ def is_eligible(config, act, arch): activation_name(config.activation) except NotImplementedError: return False - return moe_geometry(config, act, hidden_size=2 * x.shape[1]) + # SM120 kernels take whole 32-element N/K blocks; NVFP4's 4-column scale + # words then need H and I in 64s (e.g. gpt-oss, 2880). + multiple = 64 if arch == 120 else 128 + return moe_geometry(config, act, hidden_size=2 * x.shape[1], multiple=multiple) diff --git a/flashinfer/fused_moe/backends/cudnn_frost/support.py b/flashinfer/fused_moe/backends/cudnn_frost/support.py index d6acfdd1374..53f94d90b7a 100644 --- a/flashinfer/fused_moe/backends/cudnn_frost/support.py +++ b/flashinfer/fused_moe/backends/cudnn_frost/support.py @@ -13,8 +13,13 @@ def weight_view(weights, backend_key, canonical_key): return weights.get_view(key) -def moe_geometry(config, act, *, hidden_size=None): - """Cheap native-plan bounds; artifact contracts are checked by the runner.""" +def moe_geometry(config, act, *, hidden_size=None, multiple=128): + """Cheap native-plan bounds; artifact contracts are checked by the runner. + + ``multiple`` is the H/I granularity the dtype's weight layouts and kernels + need (128 unless a caller knows its layout and every packaged kernel accept + less, e.g. SM120 BF16). + """ t = act.num_tokens h = act.hidden_states_q.shape[1] if hidden_size is None else hidden_size i = config.experts.intermediate_size @@ -25,6 +30,6 @@ def moe_geometry(config, act, *, hidden_size=None): and t * k < 2**31 - 128 * e and 0 < h <= 1 << 20 and 0 < i <= 1 << 20 - and h % 128 == 0 - and i % 128 == 0 + and h % multiple == 0 + and i % multiple == 0 ) diff --git a/flashinfer/fused_moe/prepare.py b/flashinfer/fused_moe/prepare.py index e9f5dd9b86d..e14cfed5088 100644 --- a/flashinfer/fused_moe/prepare.py +++ b/flashinfer/fused_moe/prepare.py @@ -2045,6 +2045,82 @@ def quantize_experts(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: } +def _nvfp4_unswizzle_scales(scales: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Linear ``[..., rows, cols]`` bytes of 128x4-swizzled NVFP4 block scales. + + The swizzled byte of ``(row, col)`` sits at ``[row // 128, col // 4, + row % 32, row % 128 // 32, col % 4]``. + """ + padded = round_up(rows, _NVFP4_SF_SWIZZLE_ROWS) + lead = scales.shape[:-2] + return ( + scales.view(torch.uint8) + .reshape(*lead, padded // 128, cols // 4, 32, 4, 4) + .transpose(-4, -2) + .reshape(*lead, padded, cols)[..., :rows, :] + ) + + +def _nvfp4_swizzle_scales(linear: torch.Tensor) -> torch.Tensor: + """128x4-swizzle linear ``[..., rows, cols]`` scale bytes, zero-padding rows to 128.""" + *lead, rows, cols = linear.shape + padded = torch.zeros( + (*lead, round_up(rows, _NVFP4_SF_SWIZZLE_ROWS), cols), + dtype=torch.uint8, + device=linear.device, + ) + padded[..., :rows, :] = linear + return ( + padded.reshape(*lead, -1, 4, 32, cols // 4, 4) + .transpose(-4, -2) + .contiguous() + .reshape(padded.shape) + ) + + +def prepare_cudnn_frost_nvfp4_weights( + w1_bf16: torch.Tensor, + w2_bf16: torch.Tensor, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation=None, + device: Optional[torch.device] = None, +) -> Dict[str, torch.Tensor]: + """Build the ``cudnn_frost_nvfp4`` view: the CUTLASS NVFP4 view, with + per-branch FC1 block scales when the gated branches do not fill whole + 128-row scale atoms. + + The canonical view swizzles each expert's ``[up, gate]`` rows as one + matrix, while Frost reads up and gate as separate operands. When ``I`` is a + multiple of 128 the gate rows start on an atom boundary and the canonical + tensor is used as is. Otherwise (e.g. gpt-oss, ``I = 2880``) the gate half + starts mid-atom, so ``fc1_weight_block_scale`` becomes + ``uint8 [2, E, round_up(I, 128), H // 16]``: each branch swizzled on its + own, with zero padding rows. The quantized data are unchanged. + """ + view = prepare_cutlass_nvfp4_weights( + w1_bf16, + w2_bf16, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + activation=activation, + device=device, + ) + activation = _normalize_activation(activation) + if not activation.is_gated or intermediate_size % _NVFP4_SF_SWIZZLE_ROWS == 0: + return view + cols = hidden_size // _NVFP4_SF_VEC_SIZE + linear = _nvfp4_unswizzle_scales( + view["fc1_weight_block_scale"], 2 * intermediate_size, cols + ) + branches = linear.reshape(num_local_experts, 2, intermediate_size, cols) + view["fc1_weight_block_scale"] = _nvfp4_swizzle_scales(branches.transpose(0, 1)) + return view + + def _require_canonical_cutlass_bf16_weights( w1_bf16: torch.Tensor, w2_bf16: torch.Tensor, diff --git a/tests/moe/test_cudnn_frost_kernels.py b/tests/moe/test_cudnn_frost_kernels.py index 9d501b1d054..49d95bcb6a7 100644 --- a/tests/moe/test_cudnn_frost_kernels.py +++ b/tests/moe/test_cudnn_frost_kernels.py @@ -29,7 +29,11 @@ ) from flashinfer.fused_moe.backends.cudnn_frost.activations import ( ACTIVATIONS, + KNOWN_ACTIVATIONS, activation_name, + is_gated, + launch_scalars, + SWIGLU_PARAM_SCALARS, ) from flashinfer.fused_moe.backends.cudnn_frost.mxfp8 import ( runtime as mxfp8, @@ -101,6 +105,10 @@ def _graph_retention_reason(): } +# swiglu_param sources take alpha/beta/limit at launch (gpt-oss values). +_PARAM_SWIGLU = SwiGLU(alpha=1.702, beta=1.0, limit=7.0) + + def _sm120(): """SM120 kernels store through STG only: no TMA store, CTA pair or FMA.""" return torch.cuda.get_device_capability() == (12, 0) @@ -169,12 +177,16 @@ def _bf16_moe_reference(act, weights, activation=None): @graph_retention @frost_gpu -@pytest.mark.parametrize("name", [name for name in ACTIVATIONS if name != "swiglu"]) +@pytest.mark.parametrize( + "name", [name for name in KNOWN_ACTIVATIONS if name != "swiglu"] +) @pytest.mark.parametrize("tokens,topk", [(17, 1), (129, 4)]) def test_bf16_extended_activations_graph_and_routing(name, tokens, topk, monkeypatch): + if name == "swiglu_param" and not _sm120(): + pytest.skip("swiglu_param sources are packaged for SM120a") monkeypatch.setitem(sys.modules, "cudnn", None) torch.manual_seed(81) - activation = ACTIVATIONS[name]() + activation = _PARAM_SWIGLU if name == "swiglu_param" else ACTIVATIONS[name]() cfg = replace(bf16_config(topk), activation=activation) act, weights = bf16_packs(tokens=tokens, topk=topk) view = weights.get_view("cutlass_bf16") @@ -216,10 +228,11 @@ def test_bf16_extended_activations_graph_and_routing(name, tokens, topk, monkeyp assert relative_l2 < 0.005 +# A non-default SwiGLU runs the swiglu_param sources, which read alpha, beta and +# limit at launch (test_cudnn_frost_swiglu_param.py). @pytest.mark.parametrize( "activation", [ - ACTIVATIONS["swiglu"](alpha=2), ACTIVATIONS["swiglu_step"](limit=4), ACTIVATIONS["situ"](gate_scale=2), ACTIVATIONS["situ"](linear_scale=None), @@ -387,7 +400,7 @@ def test_bf16_model_geometry_artifacts_ragged_graph( @pytest.mark.parametrize( "overrides", [ - {"activation": SwiGLU(alpha=2.0)}, + {"activation": ACTIVATIONS["swiglu_step"](limit=4.0)}, {"finalize": MoEFinalizeConfig(do_finalize=False)}, {"experts": ExpertConfig(intermediate_size=256, local_expert_offset=1)}, ], @@ -826,9 +839,7 @@ def test_sm120_sources_roundtrip_without_cudnn(dtype, monkeypatch): # of one topology share a template and differ only in exact fragments. assert record["id"].startswith(f"sm120_cudnn_frost_{record['op']}_{dtype}_") stage = "fc2" if record["op"].endswith("gemm2") else "fc1" - gated = ( - stage == "fc1" and ACTIVATIONS[record["contract"]["activation"]]().is_gated - ) + gated = stage == "fc1" and is_gated(record["contract"]["activation"]) assert entry["path"] == ( f"sources/sm120_{stage}_{'gated' if gated else 'non_gated'}" f"_{'swap' if swap else 'normal'}_{'quantized' if quantized else 'bf16'}_stg.py" @@ -890,8 +901,9 @@ def test_sm120_sources_roundtrip_without_cudnn(dtype, monkeypatch): assert ast.dump(ast.parse(render_source(extracted, parameters))) == ast.dump( ast.parse(concrete) ) + # SM120 also packages swiglu_param: SwiGLU with launch-time alpha/beta/limit. assert {op for op, _, _ in families} == { - *(f"{prefix}grouped_gemm1_{name}" for name in ACTIVATIONS), + *(f"{prefix}grouped_gemm1_{name}" for name in KNOWN_ACTIVATIONS), f"{prefix}grouped_gemm2", } assert {swap for _, _, swap in families} == {False, True} @@ -1711,9 +1723,13 @@ def _assert_block_scale_geometry_roundtrip( tree ) assert {r["op"] for r in records} == { - *(f"block_scale_grouped_gemm1_{name}" for name in ACTIVATIONS), + *(f"block_scale_grouped_gemm1_{name}" for name in KNOWN_ACTIVATIONS), "block_scale_grouped_gemm2", } + # Launch-time SwiGLU parameters are packaged for SM120 only. + assert {r["arch"] for r in records if r["op"].endswith("_swiglu_param")} == { + "sm_120a" + } # Offline selection may retain different numbers of geometries and ABI # families. Every geometry of one family shares one source; # source paths must not accidentally mix activations or launch ABIs. @@ -2077,14 +2093,19 @@ def _assert_block_scale_launch_abi(runtime, *, packed, block, packed_weights=Non out = torch.empty(s, n, dtype=torch.bfloat16) workspace = torch.empty(kernel.workspace_bytes, dtype=torch.uint8) scale = torch.ones(1, 1, 1) - activation_scales = ( - { + activation_scales = {} + if kernel.activation == "situ": + activation_scales = { "gate_scale": torch.full((1, 1, 1), 4.0), "linear_scale": torch.full((1, 1, 1), 25.0), } - if kernel.activation == "situ" - else {} - ) + elif kernel.activation == "swiglu_param": + activation_scales = { + name: torch.full((1, 1, 1), value) + for name, value in zip( + SWIGLU_PARAM_SCALARS, launch_scalars(_PARAM_SWIGLU), strict=True + ) + } gemm_scales = tuple(torch.arange(groups, dtype=torch.float32) + 0.5 for _ in w) kwargs = {"gemm_scales": gemm_scales} args = runtime._launch_arguments( @@ -2205,12 +2226,14 @@ def _assert_mxfp8_grouped_kernel_and_graph(runtime, kernel, monkeypatch, mixed=F (torch.randn(e, n, k, device="cuda") * 0.1).to(x.dtype) for _ in range(2 if kernel.gated else 1) ) - if kernel.activation in ("situ", "swiglu_step"): + if kernel.activation in ("situ", "swiglu_step", "swiglu_param"): # Exercise saturation/clamps, rather than only their small-signal limits. x.copy_((x.float() * 16).to(x.dtype)) if not mixed: for weight in w: weight.copy_((weight.float() * 16).to(weight.dtype)) + # swiglu_param kernels read alpha/beta/limit at launch. + param = kernel.activation == "swiglu_param" a_sf = torch.randint(125, 128, (s, k // 32), dtype=torch.uint8, device="cuda") b_sf = tuple( torch.randint(125, 128, (e, n, k // 32), dtype=torch.uint8, device="cuda") @@ -2248,6 +2271,7 @@ def _assert_mxfp8_grouped_kernel_and_graph(runtime, kernel, monkeypatch, mixed=F out, scale=scale, gemm_scales=alphas, + **({"activation": _PARAM_SWIGLU} if param else {}), ) def reference(): @@ -2278,7 +2302,7 @@ def reference(): values = ( compute_reference_activation( values.float(), - ACTIVATIONS[kernel.activation](), + _PARAM_SWIGLU if param else ACTIVATIONS[kernel.activation](), n, out_dtype=torch.float32, ) @@ -2964,7 +2988,7 @@ def _nvfp4_dequant(data, scales): @frost_gpu -@pytest.mark.parametrize("name", [*ACTIVATIONS, "fc2"]) +@pytest.mark.parametrize("name", [*KNOWN_ACTIVATIONS, "fc2"]) @pytest.mark.parametrize( "swap_ab,store", [(False, "stg"), (False, "tma"), (True, "stg"), (True, "tma")] ) @@ -3003,7 +3027,7 @@ def test_nvfp4_grouped_kernel_extremes_alpha_and_graph( [0x00, 0x88, 0x77, 0xFF, 0x7F, 0xF7], device="cuda", dtype=torch.uint8 ) sf_values = torch.tensor([0.03125, 0.0625, 0.125], device="cuda") - if name in ("situ", "swiglu_step"): + if name in ("situ", "swiglu_step", "swiglu_param"): sf_values *= 8 a_sf = sf_values[torch.randint(3, (s, k // 16), device="cuda")].to( torch.float8_e4m3fn @@ -3023,6 +3047,8 @@ def test_nvfp4_grouped_kernel_extremes_alpha_and_graph( scale = torch.full((1, 1, 1), 0.5, device="cuda") if kernel.fc1 else None packed_a = nvfp4.pack_token_scales(a_sf, offsets) out = torch.empty(s, n, dtype=torch.bfloat16, device="cuda") + # swiglu_param kernels read alpha/beta/limit at launch. + activation = _PARAM_SWIGLU if name == "swiglu_param" else None plan = nvfp4.PreparedNvfp4GroupedGemm( kernel, x, @@ -3033,6 +3059,7 @@ def test_nvfp4_grouped_kernel_extremes_alpha_and_graph( out, scale=scale, gemm_scales=alphas, + **({"activation": activation} if activation else {}), ) def reference(): @@ -3048,7 +3075,10 @@ def reference(): if kernel.fc1: values = ( compute_reference_activation( - values, ACTIVATIONS[name](), n, out_dtype=torch.float32 + values, + activation or ACTIVATIONS[name](), + n, + out_dtype=torch.float32, ) * scale.flatten()[0] ) @@ -3800,7 +3830,7 @@ def test_mxfp8_mxfp4_auto_admission_uses_logical_hidden_size(name): @frost_gpu -@pytest.mark.parametrize("name", [*ACTIVATIONS, "fc2"]) +@pytest.mark.parametrize("name", [*KNOWN_ACTIVATIONS, "fc2"]) @pytest.mark.parametrize( "swap_ab,store", [(False, "stg"), (False, "tma"), (True, "stg"), (True, "tma")] ) diff --git a/tests/moe/test_cudnn_frost_swiglu_param.py b/tests/moe/test_cudnn_frost_swiglu_param.py new file mode 100644 index 00000000000..8b5d087f2f2 --- /dev/null +++ b/tests/moe/test_cudnn_frost_swiglu_param.py @@ -0,0 +1,318 @@ +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""cuDNN Frost SwiGLU with launch-time alpha/beta/limit (``swiglu_param``). + +DeepSeek-V4/V4.1-Flash and GLM-5.3-Flash clamp SwiGLU at 10; gpt-oss and +MiniMax-M3 use OpenAI's ``alpha=1.702, beta=1, limit=7`` form. The SM120 sources +read these values as scalar kernel inputs, so one kernel serves every model. +SM120 BF16 and NVFP4 also admit gpt-oss's H = I = 2880 (BF16 in 32-element +blocks, NVFP4 in 64s); the canonical MXFP8 weight scales still need 128s. +""" + +import pytest +import torch + +from flashinfer.fused_moe import ( + BackendOptions, + CudnnFrostBf16Config, + CutlassNvfp4Config, + CudnnFrostMxfp8Config, + CudnnFrostMxfp8Mxfp4Config, + CudnnFrostNvfp4Config, + ExecutionConfig, + ExpertConfig, + MoEActivationPack, + MoEConfig, + MoELayer, + MoEWeightPack, + QuantConfig, + QuantFormat, + RoutingConfig, + SiTU, + SwiGLU, + SwiGLUStep, +) +from flashinfer.autotuner import autotune +from flashinfer.fused_moe.api import ReLU2 +from flashinfer.fused_moe.backends.cudnn_frost.mxfp8 import support as mxfp8_support +from flashinfer.fused_moe.backends.cudnn_frost.mxfp8_mxfp4 import ( + support as mixed_support, +) +from flashinfer.fused_moe.backends.cudnn_frost.nvfp4 import support as nvfp4_support +from flashinfer.fused_moe.backends.cudnn_frost.activations import ( + SWIGLU_PARAM_SCALARS, + activation_name, + activation_scalars, + launch_scalars, +) +from flashinfer.fused_moe.backends.cudnn_frost.bf16 import moe as bf16_moe +from flashinfer.fused_moe.backends.cudnn_frost.cache import ( + require_graph_resource_retention, +) +from flashinfer.fused_moe.backends.cudnn_frost.dimension_catalog import ( + SWIGLU_PARAM_TAIL, +) +from flashinfer.fused_moe.backends.cudnn_frost.manifest import read_manifest +from flashinfer.fused_moe.backends.cudnn_frost.runtime import artifact_root +from tests.moe.test_cudnn_frost_kernels import ( + _bf16_moe_reference, + _mxfp8_moe_reference, + _nvfp4_moe_reference, +) + +CLAMPED = SwiGLU(alpha=1.0, beta=0.0, limit=10.0) # DeepSeek-V4, GLM-5.3 +OAI = SwiGLU(alpha=1.702, beta=1.0, limit=7.0) # gpt-oss, MiniMax-M3 +FORMATS = { + "bf16": (CudnnFrostBf16Config, QuantFormat.BF16, QuantFormat.BF16), + "mxfp8": (CudnnFrostMxfp8Config, QuantFormat.MXFP8, QuantFormat.MXFP8), + "nvfp4": (CudnnFrostNvfp4Config, QuantFormat.NVFP4, QuantFormat.NVFP4), + "mxfp8_mxfp4": (CudnnFrostMxfp8Mxfp4Config, QuantFormat.MXFP4, QuantFormat.MXFP8), +} + + +def _records(dtype): + path = artifact_root(dtype) / "cudnn_frost_selected_kernels.json" + return read_manifest(path)["kernels"] + + +def test_activation_names_and_launch_scalars(): + assert activation_name(SwiGLU()) == "swiglu" + assert activation_name(CLAMPED) == "swiglu_param" + assert activation_name(OAI) == "swiglu_param" + assert activation_scalars(SwiGLU()) == {} + assert activation_scalars(OAI) == dict( + swiglu_alpha=1.702, swiglu_beta=1.0, swiglu_limit=7.0, swiglu_lower=-7.0 + ) + assert launch_scalars(CLAMPED) == (1.0, 0.0, 10.0, -10.0) + assert launch_scalars(SwiGLU())[:2] == (1.0, 0.0) + # Only SwiGLU has scalar-parameterized sources. + for other in (SiTU(gate_scale=5.0), SwiGLUStep(limit=5.0)): + with pytest.raises(NotImplementedError): + activation_name(other) + + +@pytest.mark.parametrize("dtype", FORMATS) +def test_swiglu_param_catalog(dtype): + records = _records(dtype) + param = [r for r in records if r["contract"]["activation"] == "swiglu_param"] + assert param and {r["arch"] for r in param} == {"sm_120a"} + # Every SM120 SwiGLU tile/output/orientation also has its swiglu_param kernel. + + def axes(r): + return (r["tactic"]["tile"], r["contract"]["output_dtype"]) + + swiglu = { + axes(r) + for r in records + if r["arch"] == "sm_120a" and r["contract"]["activation"] == "swiglu" + } + assert {axes(r) for r in param} == swiglu + for r in param: + tail = r["launch"]["tail"] + assert set(SWIGLU_PARAM_SCALARS) <= set(tail) + scale = tail.index("scale") + assert tuple(tail[scale - len(SWIGLU_PARAM_TAIL) : scale]) == SWIGLU_PARAM_TAIL + # SM120 kernels take whole 32-element N/K blocks (e.g. 2880). + for r in records: + if r["arch"] == "sm_120a": + assert r["contract"]["n"]["multiple_of"] == 32 + assert r["contract"]["k"]["multiple_of"] == 32 + + +def test_swiglu_param_rejected_where_no_kernels_are_packaged(monkeypatch): + """SM107 packages no swiglu_param sources: the runner declines, not crashes.""" + config = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(), + experts=ExpertConfig(intermediate_size=256), + backend=BackendOptions((CudnnFrostBf16Config(),)), + activation=OAI, + ) + runner = bf16_moe.CudnnFrostBf16MoeRunner.__new__(bf16_moe.CudnnFrostBf16MoeRunner) + runner.config, runner.device = config, torch.device("cuda", 0) + monkeypatch.setattr(bf16_moe, "get_compute_capability", lambda device: (10, 7)) + with pytest.raises(NotImplementedError, match="no swiglu_param kernels"): + runner._check_support() + + +@pytest.mark.parametrize( + "dtype,eligible_on_sm120", + [("nvfp4", True), ("mxfp8", False), ("mxfp8_mxfp4", False)], +) +def test_block_scale_gpt_oss_geometry_admission(dtype, eligible_on_sm120): + """H = I = 2880: NVFP4 takes 64s on SM120; canonical MX weight scales are + int32-packed per 128 K, so MXFP8 still needs 128s. SM107 needs 128s.""" + config = MoEConfig( + routing=RoutingConfig(num_experts=32, top_k=4), + quant=QuantConfig(weight=FORMATS[dtype][1], activation=FORMATS[dtype][2]), + experts=ExpertConfig(intermediate_size=2880), + execution=ExecutionConfig(tune_max_num_tokens=64), + activation=OAI, + ) + if dtype == "nvfp4": + x, module = torch.empty(8, 1440, dtype=torch.uint8), nvfp4_support + else: + x = torch.empty(8, 2880, dtype=torch.float8_e4m3fn) + module = mxfp8_support if dtype == "mxfp8" else mixed_support + ids = torch.zeros(8, 4, dtype=torch.int32) + act = MoEActivationPack(x, None, ids, torch.full((8, 4), 0.25)) + assert module.is_eligible(config, act, 120) is eligible_on_sm120 + assert not module.is_eligible(config, act, 107) + + +def test_nvfp4_per_branch_fc1_scales(): + """I % 128 == 64: the gate rows start mid-atom in the canonical [up; gate] + scale matrix, so the Frost view swizzles each branch on its own.""" + if not torch.cuda.is_available(): + pytest.skip("fp4_quantize requires CUDA") + from flashinfer.fused_moe.prepare import ( + _nvfp4_swizzle_scales, + _nvfp4_unswizzle_scales, + ) + + e, h, i = 3, 128, 192 + gen = torch.Generator(device="cuda").manual_seed(5) + w1 = torch.randn(e, 2 * i, h, device="cuda", generator=gen).bfloat16() + w2 = torch.randn(e, h, i, device="cuda", generator=gen).bfloat16() + common = dict(num_local_experts=e, hidden_size=h, intermediate_size=i) + canonical = CutlassNvfp4Config.prepare_weights(w1, w2, **common) + frost = CudnnFrostNvfp4Config.prepare_weights(w1, w2, **common) + for name, value in canonical.items(): + if name != "fc1_weight_block_scale": + assert torch.equal(frost[name], value), name + sf = frost["fc1_weight_block_scale"] + assert sf.shape == (2, e, 256, h // 16) and sf.is_contiguous() + linear = _nvfp4_unswizzle_scales(canonical["fc1_weight_block_scale"], 2 * i, 8) + for branch in range(2): + rows = _nvfp4_unswizzle_scales(sf[branch], 256, 8) + assert torch.equal(rows[:, :i], linear[:, branch * i : (branch + 1) * i]) + assert not rows[:, i:].any() + # The pair is an exact inverse, padding aside. + assert torch.equal( + _nvfp4_swizzle_scales(linear), canonical["fc1_weight_block_scale"] + ) + # Gated FC1 with I % 128 == 0, and non-gated FC1, keep the canonical view. + for activation, size in ((SwiGLU(), 256), (ReLU2(), 192)): + rows = 2 * size if activation.is_gated else size + a = torch.randn(e, rows, h, device="cuda", generator=gen).bfloat16() + b = torch.randn(e, h, size, device="cuda", generator=gen).bfloat16() + kwargs = dict(common, intermediate_size=size, activation=activation) + assert torch.equal( + CudnnFrostNvfp4Config.prepare_weights(a, b, **kwargs)[ + "fc1_weight_block_scale" + ], + CutlassNvfp4Config.prepare_weights(a, b, **kwargs)[ + "fc1_weight_block_scale" + ], + ) + + +def _sm120_case(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0): + pytest.skip("swiglu_param sources are packaged for SM120a") + try: + require_graph_resource_retention() + except NotImplementedError as exc: + pytest.skip(str(exc)) + + +def _reference(dtype, act, weights, activation): + if dtype == "bf16": + return _bf16_moe_reference(act, weights, activation) + if dtype == "nvfp4": + return _nvfp4_moe_reference(act, weights, activation) + return _mxfp8_moe_reference(act, weights, activation, mixed=dtype == "mxfp8_mxfp4") + + +def _run(dtype, activation, geometry, tokens, seed=17): + cls, weight_format, act_format = FORMATS[dtype] + e, h, i, k = geometry + config = MoEConfig( + routing=RoutingConfig(num_experts=e, top_k=k), + quant=QuantConfig(weight=weight_format, activation=act_format), + experts=ExpertConfig(intermediate_size=i), + backend=BackendOptions((cls(),)), + activation=activation, + execution=ExecutionConfig(tune_max_num_tokens=max(tokens, 16)), + ) + layer = MoELayer(config) + gen = torch.Generator(device="cuda").manual_seed(seed) + # |gate|, |up| ~ 6: both clamps and the sigmoid knee are exercised. + scale = 3.0 / h**0.5 + w1 = torch.randn(e, 2 * i, h, device="cuda", generator=gen) * scale + w2 = torch.randn(e, h, i, device="cuda", generator=gen) * (0.5 / i**0.5) + w1, w2 = w1.bfloat16(), w2.bfloat16() + view = cls.prepare_weights( + w1, + w2, + num_local_experts=e, + hidden_size=h, + intermediate_size=i, + activation=activation, + ) + x = torch.randn(tokens, h, device="cuda", generator=gen).mul(2.0).bfloat16() + xq, xsf = cls.prepare_activations(x, quant=config.quant) + ids = torch.rand(tokens, e, device="cuda", generator=gen).topk(k).indices.int() + scores = torch.rand(tokens, k, device="cuda", generator=gen).softmax(-1) + act = MoEActivationPack(xq, xsf, ids, scores) + own = MoEWeightPack({f"cudnn_frost_{dtype}": view}) + if dtype == "nvfp4": + # The reference reads the canonical [up; gate] scale matrix. + view = CutlassNvfp4Config.prepare_weights( + w1, w2, num_local_experts=e, hidden_size=h, intermediate_size=i + ) + shared = MoEWeightPack({f"cutlass_{dtype}": view}) + with autotune(tuning_buckets=(tokens,)): + out = layer(act, own).clone() + assert layer.winner_backend == f"cudnn_frost_{dtype}" + reference = _reference(dtype, act, shared, activation) + default = _reference(dtype, act, shared, SwiGLU()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + replayed = layer(act, own) + replayed.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + return out, replayed, reference, default + + +def _relative(a, b): + return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)).item() + + +@pytest.mark.parametrize("dtype", FORMATS) +@pytest.mark.parametrize("activation", [CLAMPED, OAI], ids=["limit10", "oai"]) +@pytest.mark.parametrize("tokens", [3, 200]) +def test_swiglu_param_moe_layer(dtype, activation, tokens): + _sm120_case() + out, replayed, reference, default = _run( + dtype, activation, (16, 1024, 512, 4), tokens + ) + tolerance = 0.005 if dtype == "bf16" else 0.03 + assert torch.isfinite(out).all() + assert _relative(out, reference) < tolerance + torch.testing.assert_close(replayed, out) + # The parameters matter: the default SwiGLU result is far from this one. + assert _relative(default, reference) > 10 * tolerance + + +@pytest.mark.parametrize("dtype", ("bf16", "nvfp4")) +@pytest.mark.parametrize("tokens", [1, 37, 512]) +def test_gpt_oss_geometry(dtype, tokens): + """gpt-oss-20b: E=32, H=I=2880, top-4, SwiGLU(1.702, 1, 7) on SM120.""" + _sm120_case() + out, replayed, reference, default = _run(dtype, OAI, (32, 2880, 2880, 4), tokens) + tolerance = 0.005 if dtype == "bf16" else 0.03 + assert torch.isfinite(out).all() + assert _relative(out, reference) < tolerance + torch.testing.assert_close(replayed, out) + assert _relative(default, reference) > 10 * tolerance + + +@pytest.mark.parametrize("activation", [CLAMPED, SwiGLU()], ids=["limit10", "default"]) +def test_nvfp4_half_atom_intermediate(activation): + """H = I = 576 (4.5 atoms): per-branch FC1 scales and padded FC2 scales.""" + _sm120_case() + out, replayed, reference, _ = _run("nvfp4", activation, (8, 576, 576, 2), 77) + assert torch.isfinite(out).all() + assert _relative(out, reference) < 0.03 + torch.testing.assert_close(replayed, out)