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)