chore: import upstream snapshot with attribution

python/sglang/jit_kernel/minimax_quant_ue8m0.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_module(group_size: int) -> Module:
+    args = make_cpp_args(group_size, is_arch_support_pdl())
+    return load_jit(
+        "minimax_per_token_quant_ue8m0",
+        *args,
+        cuda_files=["minimax/per_token_quant_ue8m0.cuh"],
+        cuda_wrappers=[
+            ("per_token_quant_ue8m0", f"per_token_quant_ue8m0<{args}>"),
+        ],
+    )
+
+
+@cache_once
+def _jit_scatter_module(group_size: int, topk: int) -> Module:
+    # topk is a template arg so the dst-row load/store loops fully unroll.
+    args = make_cpp_args(group_size, topk, is_arch_support_pdl())
+    return load_jit(
+        "minimax_per_token_quant_ue8m0_scatter",
+        *args,
+        cuda_files=["minimax/per_token_quant_ue8m0.cuh"],
+        cuda_wrappers=[
+            (
+                "per_token_quant_ue8m0_scatter",
+                f"per_token_quant_ue8m0_scatter<{args}>",
+            ),
+        ],
+    )
+
+
+def per_token_quant_fp8_ue8m0(
+    x: torch.Tensor, group_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-token group quant to FP8-e4m3 with a fused UE8M0 (int32-packed) scale.
+
+    Returns ``(x_q, x_sf)`` where ``x_q`` is fp8_e4m3 ``[num_tokens, hidden]`` and
+    ``x_sf`` is the int32-packed UE8M0 scale ``[num_tokens, hidden//group_size//4]``
+    (row-major). Byte-identical to ``per_token_group_quant_fp8(scale_ue8m0=True)``
+    followed by ``transform_sf_into_required_layout`` (both ceil-round the scale),
+    but does it in a single kernel -- no separate transpose/pack launch.
+    """
+    assert x.is_cuda and x.dtype == torch.bfloat16 and x.dim() == 2
+    assert x.is_contiguous()
+    num_tokens, hidden = x.shape
+    assert hidden % group_size == 0
+    num_groups = hidden // group_size
+    assert num_groups % 4 == 0, "num_groups must be a multiple of 4 for int32 packing"
+
+    x_q = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    x_sf = torch.empty(
+        (num_tokens, num_groups // 4), dtype=torch.int32, device=x.device
+    )
+    _jit_module(group_size).per_token_quant_ue8m0(x, x_q, x_sf)
+    return x_q, x_sf
+
+
+def per_token_quant_fp8_ue8m0_scatter(
+    x: torch.Tensor,
+    gateup_input: torch.Tensor,
+    gateup_input_scale: torch.Tensor,
+    src2dst: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk: int,
+    m_max: int,
+    group_size: int = 128,
+) -> None:
+    """Fused per-token FP8/UE8M0 quant **and** scatter into the permuted grouped-GEMM
+    input -- a single kernel replacing ``per_token_quant_fp8_ue8m0`` +
+    ``fill_gateup_input_triton_kernel``.
+
+    For each source token it computes the fp8 row + int32-packed UE8M0 scale once,
+    then writes them to each of the token's ``topk`` destination rows:
+      ``gateup_input``       fp8  ``[E, m_max, hidden]``      (row ``src2dst[token, i]``)
+      ``gateup_input_scale`` int32 ``[E, hidden//group//4, m_max]`` (MN-major; byte-scattered)
+    Slots with ``topk_ids[token, i] < 0`` are skipped. Byte-identical to the
+    two-kernel path on every written row.
+    """
+    assert x.is_cuda and x.dtype == torch.bfloat16 and x.dim() == 2
+    assert x.is_contiguous()
+    assert gateup_input.dtype == torch.float8_e4m3fn and gateup_input.dim() == 3
+    assert gateup_input_scale.dtype == torch.int32 and gateup_input_scale.dim() == 3
+    num_tokens, hidden = x.shape
+    assert hidden % group_size == 0
+    num_groups = hidden // group_size
+    assert num_groups % 4 == 0, "num_groups must be a multiple of 4 for int32 packing"
+    _jit_scatter_module(group_size, int(topk)).per_token_quant_ue8m0_scatter(
+        x, gateup_input, gateup_input_scale, src2dst, topk_ids, int(topk), int(m_max)
+    )