sglang 0.4.5__py3-none-any.whl → 0.4.5.post2__py3-none-any.whl

sglang/test/test_block_fp8.py

@@ -7,10 +7,13 @@ import torch
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
 from sglang.srt.layers.quantization.fp8_kernel import (
+    per_tensor_quant_mla_deep_gemm_masked_fp8,
+    per_tensor_quant_mla_fp8,
     per_token_group_quant_fp8,
     static_quant_fp8,
     w8a8_block_fp8_matmul,
 )
+from sglang.srt.layers.quantization.fp8_utils import input_to_float8
 from sglang.test.test_utils import CustomTestCase
 
 _is_cuda = torch.cuda.is_available() and torch.version.cuda
@@ -155,6 +158,117 @@ class TestStaticQuantFP8(CustomTestCase):
                 self._static_quant_fp8(*params)
 
 
+class TestPerTensorQuantMlaFP8(CustomTestCase):
+    DTYPES = [torch.half, torch.bfloat16, torch.float32]
+    NUM_TOKENS = [7, 83, 2048]
+    D = [512, 4096, 5120, 13824]
+    LAST_D_EXT = [1024, 0]
+    LAST_D = [512]
+    SEEDS = [0]
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        torch.set_default_device("cuda")
+
+    def _per_tensor_quant_mla_fp8(self, num_tokens, d, last_d_ext, last_d, dtype, seed):
+        torch.manual_seed(seed)
+
+        x = torch.rand(
+            (num_tokens, d // last_d, last_d + last_d_ext),
+            dtype=dtype,
+        )
+        x_sub, _ = x.split([last_d, last_d_ext], dim=-1)
+
+        with torch.inference_mode():
+            ref_out, ref_s = input_to_float8(x_sub.transpose(0, 1))
+            out, out_s = per_tensor_quant_mla_fp8(x_sub.transpose(0, 1))
+
+        self.assertTrue(out.is_contiguous())
+        self.assertTrue(
+            torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.50)
+        )
+        self.assertTrue(
+            torch.allclose(out_s.to(torch.float32), ref_s.to(torch.float32))
+        )
+
+    def test_per_tensor_quant_mla_fp8(self):
+        for params in itertools.product(
+            self.NUM_TOKENS,
+            self.D,
+            self.LAST_D_EXT,
+            self.LAST_D,
+            self.DTYPES,
+            self.SEEDS,
+        ):
+            with self.subTest(
+                num_tokens=params[0],
+                d=params[1],
+                last_d_ext=params[2],
+                last_d=params[3],
+                dtype=params[4],
+                seed=params[5],
+            ):
+                self._per_tensor_quant_mla_fp8(*params)
+
+
+class TestPerTokenGroupQuantMlaDeepGemmMaskedFP8(CustomTestCase):
+    DTYPES = [torch.half, torch.bfloat16, torch.float32]
+    B = [128]
+    NUM_TOKENS = [7, 83, 2048, 1024 * 16]
+    D = [512, 128]
+    GROUP_SIZE = [128]
+    SEEDS = [0]
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        torch.set_default_device("cuda")
+
+    def _per_token_group_quant_mla_deep_gemm_masked_fp8(
+        self, b, num_tokens, d, dtype, group_size, seed
+    ):
+        torch.manual_seed(seed)
+
+        x = torch.rand(b, num_tokens, d, dtype=dtype)
+
+        with torch.inference_mode():
+            ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size, 1e-12)
+            out, scale, _, _, _ = per_tensor_quant_mla_deep_gemm_masked_fp8(
+                x, group_size
+            )
+            out = out[:, :num_tokens, :]
+            scale = scale[:, :num_tokens, :]
+
+        self.assertTrue(
+            torch.allclose(
+                out.to(torch.float32), ref_out.to(torch.float32), rtol=0.20, atol=1e-2
+            )
+        )
+        self.assertTrue(torch.allclose(scale, ref_scale))
+
+    def test_per_token_group_quant_mla_deep_gemm_masked_fp8(self):
+        for params in itertools.product(
+            self.B,
+            self.NUM_TOKENS,
+            self.D,
+            self.DTYPES,
+            self.GROUP_SIZE,
+            self.SEEDS,
+        ):
+            with self.subTest(
+                b=params[0],
+                num_tokens=params[1],
+                d=params[2],
+                dtype=params[3],
+                group_size=params[4],
+                seed=params[5],
+            ):
+                self._per_token_group_quant_mla_deep_gemm_masked_fp8(*params)
+
+
 # For test
 def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.float16):
     """This function performs matrix multiplication with block-wise quantization using native torch.
@@ -428,5 +542,115 @@ class TestW8A8BlockFP8FusedMoE(CustomTestCase):
                 self._w8a8_block_fp8_fused_moe(*params)
 
 
+# For test
+def torch_w8a8_block_fp8_bmm(a, a_s, w, w_s, block_shape, out_dtype):
+    """This function performs bmm with block-wise quantization using native torch."""
+
+    B, N, _ = w.shape
+    _, M, _ = a.shape
+    out = torch.empty((B, M, N), dtype=out_dtype, device=a.device)
+
+    for i in range(B):
+        out[i] = native_w8a8_block_fp8_matmul(
+            a[i], w[i], a_s[i], w_s[i], block_shape, output_dtype=out_dtype
+        )
+
+    return out
+
+
+class TestW8A8BlockFP8BatchedDeepGemm(CustomTestCase):
+    DTYPES = [torch.bfloat16]
+    M = [1, 33, 64, 222, 8192]
+    N = [128, 512]
+    K = [128, 512]
+    BATCH = [128]
+    BLOCK_SIZE = [[128, 128]]
+    SEEDS = [0]
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        try:
+            import deep_gemm
+        except ImportError:
+            raise unittest.SkipTest("DeepGEMM is not available")
+        torch.set_default_device("cuda")
+
+    def _w8a8_block_fp8_batched_deep_gemm(self, M, N, K, B, block_size, dtype, seed):
+        torch.manual_seed(seed)
+        factor_for_scale = 1e-2
+        fp8_info = torch.finfo(torch.float8_e4m3fn)
+        fp8_max, fp8_min = fp8_info.max, fp8_info.min
+
+        a_fp32 = torch.randn((B, M, K), dtype=torch.float32) / 10
+        a = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+        w_fp32 = (torch.rand((B, N, K), dtype=torch.float32) - 0.5) * 2 * fp8_max
+        w = w_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+        block_n, block_k = block_size[0], block_size[1]
+        n_tiles_w = (N + block_n - 1) // block_n
+        k_tiles_w = (K + block_k - 1) // block_k
+
+        w_s = (
+            torch.rand((B, n_tiles_w, k_tiles_w), dtype=torch.float32)
+            * factor_for_scale
+        )
+        a_s = torch.rand((B, M, k_tiles_w), dtype=torch.float32) * factor_for_scale
+
+        ae = a.new_empty(B, (M + 255) // 256 * 256, K)
+        ae_s = a_s.new_empty(B, (M + 255) // 256 * 256, k_tiles_w)
+        oe = torch.empty((B, (M + 255) // 256 * 256, N), dtype=dtype)
+        ae[:, :M, :] = a
+        ae_s[:, :M, :] = a_s
+
+        masked_m = torch.full((B,), M, dtype=torch.int)
+        expected_m = M
+        lhs = (
+            ae,
+            ae_s,
+        )
+        rhs = (
+            w,
+            w_s,
+        )
+
+        from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_masked
+
+        with torch.inference_mode():
+            ref_out = torch_w8a8_block_fp8_bmm(a, a_s, w, w_s, block_size, dtype)
+            m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs, rhs, oe, masked_m, expected_m)
+            out = oe[:, :M, :]
+
+        self.assertTrue(
+            torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
+            / torch.mean(torch.abs(ref_out.to(torch.float32)))
+            < 0.0001
+        )
+
+    def test_w8a8_block_fp8_batched_deep_gemm(self):
+
+        for params in itertools.product(
+            self.M,
+            self.N,
+            self.K,
+            self.BATCH,
+            self.BLOCK_SIZE,
+            self.DTYPES,
+            self.SEEDS,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                B=params[3],
+                block_size=params[4],
+                dtype=params[5],
+                seed=params[6],
+            ):
+                self._w8a8_block_fp8_batched_deep_gemm(*params)
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)

sglang/test/test_custom_ops.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-from sglang.srt.custom_op import scaled_fp8_quant
+from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.utils import is_cuda
 
 

sglang/test/test_utils.py

@@ -25,7 +25,12 @@ from sglang.bench_serving import run_benchmark
 from sglang.global_config import global_config
 from sglang.lang.backend.openai import OpenAI
 from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
-from sglang.srt.utils import get_bool_env_var, kill_process_tree, retry
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_port_available,
+    kill_process_tree,
+    retry,
+)
 from sglang.test.run_eval import run_eval
 from sglang.utils import get_exception_traceback
 
@@ -37,11 +42,6 @@ DEFAULT_FP8_MODEL_NAME_FOR_DYNAMIC_QUANT_ACCURACY_TEST = (
 DEFAULT_FP8_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST = (
     "nvidia/Llama-3.1-8B-Instruct-FP8"
 )
-# TODO(yundai424): right now specifying to an older revision since the latest one
-#  carries kv cache quantization which doesn't work yet
-DEFAULT_FP8_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_REVISION = (
-    "13858565416dbdc0b4e7a4a677fadfbd5b9e5bb9"
-)
 
 DEFAULT_MODEL_NAME_FOR_TEST = "meta-llama/Llama-3.1-8B-Instruct"
 DEFAULT_SMALL_MODEL_NAME_FOR_TEST = "meta-llama/Llama-3.2-1B-Instruct"
@@ -103,6 +103,17 @@ def call_generate_lightllm(prompt, temperature, max_tokens, stop=None, url=None)
     return pred
 
 
+def find_available_port(base_port: int):
+    port = base_port + random.randint(100, 1000)
+    while True:
+        if is_port_available(port):
+            return port
+        if port < 60000:
+            port += 42
+        else:
+            port -= 43
+
+
 def call_generate_vllm(prompt, temperature, max_tokens, stop=None, n=1, url=None):
     assert url is not None
 
@@ -674,8 +685,6 @@ def run_bench_one_batch(model, other_args):
         "python3",
         "-m",
         "sglang.bench_one_batch",
-        "--model-path",
-        model,
         "--batch-size",
         "1",
         "--input",
@@ -684,6 +693,8 @@ def run_bench_one_batch(model, other_args):
         "8",
         *[str(x) for x in other_args],
     ]
+    if model is not None:
+        command += ["--model-path", model]
     process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
     try:

sglang/version.py

@@ -1 +1 @@
-__version__ = "0.4.5"
+__version__ = "0.4.5.post2"

{sglang-0.4.5.dist-info → sglang-0.4.5.post2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sglang
-Version: 0.4.5
+Version: 0.4.5.post2
 Summary: SGLang is yet another fast serving framework for large language models and vision language models.
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -239,20 +239,30 @@ Requires-Dist: python-multipart; extra == "runtime-common"
 Requires-Dist: pyzmq>=25.1.2; extra == "runtime-common"
 Requires-Dist: soundfile==0.13.1; extra == "runtime-common"
 Requires-Dist: torchao>=0.7.0; extra == "runtime-common"
-Requires-Dist: transformers==4.51.0; extra == "runtime-common"
+Requires-Dist: transformers==4.51.1; extra == "runtime-common"
 Requires-Dist: uvicorn; extra == "runtime-common"
 Requires-Dist: uvloop; extra == "runtime-common"
 Requires-Dist: compressed-tensors; extra == "runtime-common"
 Requires-Dist: xgrammar==0.1.17; extra == "runtime-common"
 Provides-Extra: srt
 Requires-Dist: sglang[runtime_common]; extra == "srt"
-Requires-Dist: sgl-kernel==0.0.8; extra == "srt"
+Requires-Dist: sgl-kernel==0.0.9.post2; extra == "srt"
 Requires-Dist: flashinfer_python==0.2.3; extra == "srt"
 Requires-Dist: torch==2.5.1; extra == "srt"
+Requires-Dist: torchvision==0.20.1; extra == "srt"
 Requires-Dist: cuda-python; extra == "srt"
 Requires-Dist: outlines<=0.1.11,>=0.0.44; extra == "srt"
 Requires-Dist: partial_json_parser; extra == "srt"
 Requires-Dist: einops; extra == "srt"
+Provides-Extra: blackwell
+Requires-Dist: sglang[runtime_common]; extra == "blackwell"
+Requires-Dist: sgl-kernel; extra == "blackwell"
+Requires-Dist: torch; extra == "blackwell"
+Requires-Dist: torchvision; extra == "blackwell"
+Requires-Dist: cuda-python; extra == "blackwell"
+Requires-Dist: outlines<=0.1.11,>=0.0.44; extra == "blackwell"
+Requires-Dist: partial_json_parser; extra == "blackwell"
+Requires-Dist: einops; extra == "blackwell"
 Provides-Extra: srt-hip
 Requires-Dist: sglang[runtime_common]; extra == "srt-hip"
 Requires-Dist: torch; extra == "srt-hip"
@@ -371,7 +381,7 @@ SGLang is a fast serving framework for large language models and vision language
 It makes your interaction with models faster and more controllable by co-designing the backend runtime and frontend language.
 The core features include:
 
-- **Fast Backend Runtime**: Provides efficient serving with RadixAttention for prefix caching, zero-overhead CPU scheduler, continuous batching, token attention (paged attention), speculative decoding, tensor parallelism, chunked prefill, structured outputs, and quantization (FP8/INT4/AWQ/GPTQ).
+- **Fast Backend Runtime**: Provides efficient serving with RadixAttention for prefix caching, zero-overhead CPU scheduler, continuous batching, token attention (paged attention), speculative decoding, tensor parallelism, chunked prefill, structured outputs, quantization (FP8/INT4/AWQ/GPTQ), and multi-lora batching.
 - **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, including chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions.
 - **Extensive Model Support**: Supports a wide range of generative models (Llama, Gemma, Mistral, QWen, DeepSeek, LLaVA, etc.), embedding models (e5-mistral, gte, mcdse) and reward models (Skywork), with easy extensibility for integrating new models.
 - **Active Community**: SGLang is open-source and backed by an active community with industry adoption.
@@ -391,7 +401,7 @@ Learn more in the release blogs: [v0.2 blog](https://lmsys.org/blog/2024-07-25-s
 
 ## Adoption and Sponsorship
 The project has been deployed to large-scale production, generating trillions of tokens every day.
-It is supported by the following institutions: AMD, Atlas Cloud, Baseten, Cursor, DataCrunch, Etched, Hyperbolic, Iflytek, Jam & Tea Studios, LinkedIn, LMSYS, Meituan, Nebius, Novita AI, NVIDIA, RunPod, Stanford, UC Berkeley, UCLA, xAI, and 01.AI.
+It is supported by the following institutions: AMD, Atlas Cloud, Baseten, Cursor, DataCrunch, Etched, Hyperbolic, Iflytek, Jam & Tea Studios, LinkedIn, LMSYS, Meituan, Nebius, Novita AI, NVIDIA, Oracle, RunPod, Stanford, UC Berkeley, UCLA, xAI, and 01.AI.
 
 <img src="https://raw.githubusercontent.com/sgl-project/sgl-learning-materials/main/slides/adoption.png" alt="logo" width="800" margin="10px"></img>