sglang 0.5.3rc2__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/utils/profile_merger.py

@@ -0,0 +1,199 @@
+"""Merge Chrome trace files from multiple ranks (TP, DP, PP, EP) into a single trace."""
+
+import glob
+import gzip
+import json
+import logging
+import os
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class ProfileMerger:
+    """Merge profile traces from all parallelism types: TP, DP, PP, EP."""
+
+    def __init__(self, output_dir: str, profile_id: str):
+        self.output_dir = output_dir
+        self.profile_id = profile_id
+        self.merged_trace_path = os.path.join(
+            output_dir, f"merged-{profile_id}.trace.json.gz"
+        )
+
+        # Rank types in priority order (used for sorting and labeling)
+        self.rank_types = ["tp", "dp", "pp", "ep"]
+
+        # Sort index multipliers: DP (highest) > EP > PP > TP (lowest)
+        # These ensure proper visual ordering in trace viewer
+        self.sort_index_multipliers = {
+            "dp_rank": 100_000_000,
+            "ep_rank": 1_000_000,
+            "pp_rank": 10_000,
+            "tp_rank": 100,
+        }
+
+        # PID threshold for sort_index updates (only update for system PIDs < 1000)
+        self.pid_sort_index_threshold = 1000
+
+    def merge_chrome_traces(self) -> str:
+        """Merge Chrome traces from all ranks into a single trace.
+
+        Returns:
+            Path to merged trace file.
+
+        Raises:
+            ValueError: If no trace files found.
+        """
+        trace_files = self._discover_trace_files()
+        if not trace_files:
+            raise ValueError(f"No trace files found for profile_id: {self.profile_id}")
+
+        logger.info(f"Found {len(trace_files)} trace files to merge")
+
+        merged_trace = {"traceEvents": []}
+        all_device_properties = []
+
+        for trace_file in sorted(trace_files, key=self._get_rank_sort_key):
+            rank_info = self._extract_rank_info(trace_file)
+            logger.info(f"Processing {trace_file} with rank info: {rank_info}")
+
+            output = self._handle_file(trace_file, rank_info)
+
+            merged_trace["traceEvents"].extend(output["traceEvents"])
+
+            if "deviceProperties" in output:
+                all_device_properties.extend(output["deviceProperties"])
+                del output["deviceProperties"]
+
+            for key, value in output.items():
+                if key != "traceEvents" and key not in merged_trace:
+                    merged_trace[key] = value
+
+        if all_device_properties:
+            merged_trace["deviceProperties"] = all_device_properties
+
+        with gzip.open(self.merged_trace_path, "wb") as f:
+            f.write(json.dumps(merged_trace).encode("utf-8"))
+
+        logger.info(f"Merged profile saved to: {self.merged_trace_path}")
+        logger.info(f"Total events merged: {len(merged_trace['traceEvents'])}")
+
+        return self.merged_trace_path
+
+    def _discover_trace_files(self) -> List[str]:
+        """Discover trace files matching profile_id (supports TP/DP/PP/EP formats)."""
+        patterns = [f"{self.profile_id}*.trace.json.gz"]
+
+        trace_files = []
+        for pattern in patterns:
+            search_pattern = os.path.join(self.output_dir, pattern)
+            trace_files.extend(glob.glob(search_pattern))
+
+        trace_files = [
+            f
+            for f in trace_files
+            if not f.endswith(f"merged-{self.profile_id}.trace.json.gz")
+            and not f.endswith("-memory.pickle")
+            and "TP-" in f
+        ]
+        trace_files = list(set(trace_files))
+        return trace_files
+
+    def _extract_rank_info(self, filename: str) -> Dict[str, int]:
+        """Extract rank info (TP/DP/PP/EP) from filename."""
+        basename = os.path.basename(filename)
+        rank_info = {}
+
+        for rank_type in self.rank_types:
+            match = re.search(rf"{rank_type.upper()}-(\d+)", basename)
+            if match:
+                rank_info[f"{rank_type}_rank"] = int(match.group(1))
+
+        return rank_info
+
+    def _create_rank_label(self, rank_info: Dict[str, int]) -> str:
+        parts = []
+        for rank_type in self.rank_types:
+            rank_key = f"{rank_type}_rank"
+            if rank_key in rank_info:
+                parts.append(f"{rank_type.upper()}{rank_info[rank_key]:02d}")
+
+        return f"[{'-'.join(parts)}]" if parts else "[Unknown]"
+
+    def _handle_file(self, path: str, rank_info: Dict[str, int]) -> Dict[str, Any]:
+        logger.info(f"Processing file: {path}")
+
+        try:
+            with gzip.open(path, "rt", encoding="utf-8") as f:
+                trace = json.load(f)
+
+            output = {
+                key: value for key, value in trace.items() if key != "traceEvents"
+            }
+            output["traceEvents"] = self._process_events(
+                trace.get("traceEvents", []), rank_info
+            )
+            return output
+
+        except Exception as e:
+            logger.error(f"Failed to process trace file {path}: {e}")
+            return {"traceEvents": []}
+
+    def _process_events(
+        self, events: List[Dict], rank_info: Dict[str, int]
+    ) -> List[Dict]:
+        """Process events: update sort_index and add rank labels to PIDs."""
+        rank_label = self._create_rank_label(rank_info)
+
+        for event in events:
+            if event.get("name") == "process_sort_index":
+                pid = self._maybe_cast_int(event.get("pid"))
+                if pid is not None and pid < self.pid_sort_index_threshold:
+                    event["args"]["sort_index"] = self._calculate_sort_index(
+                        rank_info, pid
+                    )
+
+            event["pid"] = f"{rank_label} {event['pid']}"
+
+        return events
+
+    def _calculate_sort_index(self, rank_info: Dict[str, int], pid: int) -> int:
+        sort_index = pid
+        for rank_type, multiplier in self.sort_index_multipliers.items():
+            sort_index += rank_info.get(rank_type, 0) * multiplier
+        return sort_index
+
+    def _get_rank_sort_key(self, path: str) -> Tuple[int, int, int, int]:
+        rank_info = self._extract_rank_info(path)
+        return tuple(
+            rank_info.get(f"{rank_type}_rank", 0)
+            for rank_type in ["dp", "ep", "pp", "tp"]
+        )
+
+    def _maybe_cast_int(self, x) -> Optional[int]:
+        try:
+            return int(x)
+        except (ValueError, TypeError):
+            return None
+
+    def get_merge_summary(self) -> Dict[str, Any]:
+        if not os.path.exists(self.merged_trace_path):
+            return {"error": "Merged trace file not found"}
+
+        try:
+            with gzip.open(self.merged_trace_path, "rt") as f:
+                merged_data = json.load(f)
+
+            trace_files = self._discover_trace_files()
+
+            return {
+                "merged_file": self.merged_trace_path,
+                "total_events": len(merged_data.get("traceEvents", [])),
+                "total_files": len(trace_files),
+                "source_files": [os.path.basename(f) for f in trace_files],
+                "profile_id": self.profile_id,
+                "device_properties_count": len(merged_data.get("deviceProperties", [])),
+            }
+        except Exception as e:
+            return {"error": f"Failed to read merged trace: {str(e)}"}

sglang/test/attention/test_flashattn_backend.py

@@ -66,7 +66,7 @@ class MockModelRunner:
             enable_memory_saver=False,
         )
         # Required by torch native backend
-        self.server_args = ServerArgs(model_path="fake_model_path")
+        self.server_args = ServerArgs(model_path="dummy")
 
 
 @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")

sglang/test/attention/test_flashattn_mla_backend.py

@@ -4,7 +4,6 @@ import torch
 
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
-from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode

sglang/test/attention/test_prefix_chunk_info.py

@@ -2,8 +2,6 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
-from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.test.test_utils import CustomTestCase

sglang/test/attention/test_trtllm_mla_backend.py

@@ -16,10 +16,15 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLABackend,
     TRTLLMMLADecodeMetadata,
 )
-from sglang.srt.layers.attention.utils import TRITON_PAD_NUM_PAGE_PER_BLOCK
+from sglang.srt.layers.attention.utils import get_num_page_per_block_flashmla
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.server_args import (
+    ServerArgs,
+    get_global_server_args,
+    set_global_server_args_for_scheduler,
+)
 from sglang.srt.utils import is_flashinfer_available
 from sglang.test.test_utils import CustomTestCase
 
@@ -104,15 +109,15 @@ TEST_CASES = {
             "page_size": 32,
             "description": "Single FP16 vs reference",
         },
-        {
-            "name": "single_fp8",
-            "batch_size": 1,
-            "max_seq_len": 64,
-            "page_size": 64,
-            "tolerance": 1e-1,
-            "kv_cache_dtype": torch.float8_e4m3fn,
-            "description": "Single FP8 vs reference",
-        },
+        # {
+        #     "name": "single_fp8",
+        #     "batch_size": 1,
+        #     "max_seq_len": 64,
+        #     "page_size": 64,
+        #     "tolerance": 1e-1,
+        #     "kv_cache_dtype": torch.float8_e4m3fn,
+        #     "description": "Single FP8 vs reference",
+        # },
         {
             "name": "batch_fp16",
             "batch_size": 32,
@@ -120,15 +125,15 @@ TEST_CASES = {
             "page_size": 32,
             "description": "Batch FP16 vs reference",
         },
-        {
-            "name": "batch_fp8",
-            "batch_size": 32,
-            "max_seq_len": 64,
-            "page_size": 64,
-            "tolerance": 1e-1,
-            "kv_cache_dtype": torch.float8_e4m3fn,
-            "description": "Batch FP8 vs reference",
-        },
+        # {
+        #     "name": "batch_fp8",
+        #     "batch_size": 32,
+        #     "max_seq_len": 64,
+        #     "page_size": 64,
+        #     "tolerance": 1e-1,
+        #     "kv_cache_dtype": torch.float8_e4m3fn,
+        #     "description": "Batch FP8 vs reference",
+        # },
     ],
     "page_size_consistency": [
         # Only 32 and 64 supported for now in flashinfer TRTLLM-GEN MLA kernel
@@ -213,13 +218,7 @@ class MockModelRunner:
         self.page_size = config["page_size"]
 
         # Server args stub - needed by attention backends
-        self.server_args = type(
-            "ServerArgs",
-            (),
-            {
-                "enable_dp_attention": False,  # Default value for testing
-            },
-        )
+        self.server_args = get_global_server_args()
 
         # Model-config stub with MLA attributes
         self.model_config = type(
@@ -320,6 +319,17 @@ def compare_outputs(trtllm_out, reference_out, tolerance=1e-2):
 class TestTRTLLMMLA(CustomTestCase):
     """Test suite for TRTLLM MLA backend with centralized configuration."""
 
+    @classmethod
+    def setUpClass(cls):
+        """Set up global server args for testing."""
+        server_args = ServerArgs(model_path="dummy")
+        server_args.enable_dp_attention = False
+        set_global_server_args_for_scheduler(server_args)
+
+    @classmethod
+    def tearDownClass(cls):
+        pass
+
     def _merge_config(self, test_case):
         """Merge test case with default configuration."""
         config = DEFAULT_CONFIG.copy()
@@ -841,25 +851,17 @@ class TestTRTLLMMLA(CustomTestCase):
                 backend.init_forward_metadata(fb)
 
                 # Verify metadata exists
-                self.assertIsNotNone(backend.forward_metadata)
-                self.assertIsInstance(backend.forward_metadata, TRTLLMMLADecodeMetadata)
+                self.assertIsNotNone(backend.forward_decode_metadata)
+                self.assertIsInstance(
+                    backend.forward_decode_metadata, TRTLLMMLADecodeMetadata
+                )
 
                 # Test metadata structure
-                metadata = backend.forward_metadata
-                self.assertIsNotNone(
-                    metadata.workspace, "Workspace should be allocated"
-                )
+                metadata = backend.forward_decode_metadata
                 self.assertIsNotNone(
                     metadata.block_kv_indices, "Block KV indices should be created"
                 )
 
-                # Test workspace properties
-                self.assertEqual(metadata.workspace.device.type, "cuda")
-                self.assertEqual(metadata.workspace.dtype, torch.uint8)
-                self.assertGreater(
-                    metadata.workspace.numel(), 0, "Workspace should have non-zero size"
-                )
-
                 # Test block KV indices properties
                 self.assertEqual(metadata.block_kv_indices.device.type, "cuda")
                 self.assertEqual(metadata.block_kv_indices.dtype, torch.int32)
@@ -915,9 +917,10 @@ class TestTRTLLMMLA(CustomTestCase):
 
                 # Should satisfy TRT-LLM and Triton constraints
                 trtllm_constraint = 128 // scenario["page_size"]
-                constraint_lcm = math.lcm(
-                    trtllm_constraint, TRITON_PAD_NUM_PAGE_PER_BLOCK
+                triton_constraint = get_num_page_per_block_flashmla(
+                    scenario["page_size"]
                 )
+                constraint_lcm = math.lcm(trtllm_constraint, triton_constraint)
                 self.assertEqual(
                     calculated_blocks % constraint_lcm,
                     0,
@@ -965,7 +968,7 @@ class TestTRTLLMMLA(CustomTestCase):
 
                 # Initialize metadata
                 backend.init_forward_metadata(fb)
-                metadata = backend.forward_metadata
+                metadata = backend.forward_decode_metadata
 
                 # Verify KV indices structure
                 block_kv_indices = metadata.block_kv_indices
@@ -1016,7 +1019,6 @@ class TestTRTLLMMLA(CustomTestCase):
 
         # Verify CUDA graph buffers are allocated
         self.assertIsNotNone(backend.decode_cuda_graph_kv_indices)
-        self.assertIsNotNone(backend.decode_cuda_graph_workspace)
 
         # Test capture metadata
         seq_lens = torch.full(
@@ -1038,7 +1040,6 @@ class TestTRTLLMMLA(CustomTestCase):
         self.assertIn(batch_size, backend.decode_cuda_graph_metadata)
         capture_metadata = backend.decode_cuda_graph_metadata[batch_size]
 
-        self.assertIsNotNone(capture_metadata.workspace)
         self.assertIsNotNone(capture_metadata.block_kv_indices)
 
         # Test replay with different sequence lengths
@@ -1061,11 +1062,8 @@ class TestTRTLLMMLA(CustomTestCase):
         )
 
         # Verify replay updated the metadata
-        replay_metadata = backend.forward_metadata
+        replay_metadata = backend.forward_decode_metadata
         self.assertIsNotNone(replay_metadata)
-        self.assertEqual(
-            replay_metadata.workspace.data_ptr(), capture_metadata.workspace.data_ptr()
-        )
 
     def test_metadata_consistency_across_calls(self):
         """Test metadata consistency across multiple forward calls."""
@@ -1083,7 +1081,7 @@ class TestTRTLLMMLA(CustomTestCase):
             config["batch_size"], seq_lens_1, backend, model_runner, config
         )
         backend.init_forward_metadata(fb_1)
-        metadata_1 = backend.forward_metadata
+        metadata_1 = backend.forward_decode_metadata
 
         # Second call with same sequence lengths
         seq_lens_2 = torch.tensor([32, 48], device=config["device"])
@@ -1091,10 +1089,9 @@ class TestTRTLLMMLA(CustomTestCase):
             config["batch_size"], seq_lens_2, backend, model_runner, config
         )
         backend.init_forward_metadata(fb_2)
-        metadata_2 = backend.forward_metadata
+        metadata_2 = backend.forward_decode_metadata
 
         # Metadata structure should be consistent
-        self.assertEqual(metadata_1.workspace.shape, metadata_2.workspace.shape)
         self.assertEqual(
             metadata_1.block_kv_indices.shape, metadata_2.block_kv_indices.shape
         )
@@ -1105,10 +1102,9 @@ class TestTRTLLMMLA(CustomTestCase):
             config["batch_size"], seq_lens_3, backend, model_runner, config
         )
         backend.init_forward_metadata(fb_3)
-        metadata_3 = backend.forward_metadata
+        metadata_3 = backend.forward_decode_metadata
 
         # Should still have valid structure
-        self.assertIsNotNone(metadata_3.workspace)
         self.assertIsNotNone(metadata_3.block_kv_indices)
         self.assertEqual(metadata_3.block_kv_indices.shape[0], config["batch_size"])
 
@@ -1263,6 +1259,178 @@ class TestTRTLLMMLA(CustomTestCase):
                     f"Max diff: {(out_trtllm - out_reference).abs().max().item()}",
                 )
 
+    def test_draft_extend_padding_unpadding_kernels(self):
+        """Test TRTLLM MLA Triton kernels: pad_draft_extend_query_kernel and unpad_draft_extend_output_kernel."""
+
+        # Import the kernels
+        from sglang.srt.layers.attention.trtllm_mla_backend import (
+            pad_draft_extend_query_kernel,
+            unpad_draft_extend_output_kernel,
+        )
+
+        def _create_test_data(
+            self, batch_size, max_seq_len, num_heads, head_dim, dtype=torch.float32
+        ):
+            """Create test data for kernel testing."""
+            device = torch.device("cuda")
+
+            # Create sequence lengths (varying lengths for each batch)
+            seq_lens = torch.randint(
+                1, max_seq_len + 1, (batch_size,), device=device, dtype=torch.int32
+            )
+
+            # Create cumulative sequence lengths
+            cum_seq_lens = torch.zeros(batch_size + 1, device=device, dtype=torch.int32)
+            cum_seq_lens[1:] = torch.cumsum(seq_lens, dim=0)
+
+            # Create input query tensor (flattened format)
+            total_tokens = cum_seq_lens[-1].item()
+            q_input = torch.randn(
+                total_tokens, num_heads, head_dim, device=device, dtype=dtype
+            )
+
+            # Create padded query tensor (batch format)
+            padded_q = torch.zeros(
+                batch_size, max_seq_len, num_heads, head_dim, device=device, dtype=dtype
+            )
+
+            return q_input, padded_q, seq_lens, cum_seq_lens
+
+        def _create_test_output_data(
+            self,
+            batch_size,
+            token_per_batch,
+            tp_q_head_num,
+            v_head_dim,
+            dtype=torch.float32,
+        ):
+            """Create test data for unpad kernel testing."""
+            device = torch.device("cuda")
+
+            # Create accept lengths (varying lengths for each batch)
+            accept_lengths = torch.randint(
+                1, token_per_batch + 1, (batch_size,), device=device, dtype=torch.int32
+            )
+
+            # Create cumulative accept lengths
+            cum_accept_lengths = torch.zeros(
+                batch_size + 1, device=device, dtype=torch.int32
+            )
+            cum_accept_lengths[1:] = torch.cumsum(accept_lengths, dim=0)
+
+            # Create raw output tensor (batch format)
+            raw_out = torch.randn(
+                batch_size,
+                token_per_batch,
+                tp_q_head_num,
+                v_head_dim,
+                device=device,
+                dtype=dtype,
+            )
+
+            # Create output tensor (flattened format)
+            total_tokens = cum_accept_lengths[-1].item()
+            output = torch.empty(
+                total_tokens, tp_q_head_num, v_head_dim, device=device, dtype=dtype
+            )
+
+            return raw_out, output, accept_lengths, cum_accept_lengths
+
+        # Test 1: pad_draft_extend_query_kernel basic functionality
+        with self.subTest(test="pad_kernel_basic"):
+            batch_size = 4
+            max_seq_len = 8
+            num_heads = 16
+            head_dim = 64
+
+            q_input, padded_q, seq_lens, cum_seq_lens = _create_test_data(
+                self, batch_size, max_seq_len, num_heads, head_dim
+            )
+
+            # Launch kernel
+            BLOCK_SIZE = 64
+            grid = (batch_size * max_seq_len,)
+
+            pad_draft_extend_query_kernel[grid](
+                q_ptr=q_input,
+                padded_q_ptr=padded_q,
+                seq_lens_q_ptr=seq_lens,
+                cumsum_ptr=cum_seq_lens,
+                batch_size=batch_size,
+                max_seq_len=max_seq_len,
+                num_heads=num_heads,
+                head_dim=head_dim,
+                BLOCK_SIZE=BLOCK_SIZE,
+            )
+
+            # Verify the padding worked correctly
+            for i in range(batch_size):
+                seq_len = seq_lens[i].item()
+
+                # Check that valid positions are copied correctly
+                for pos in range(seq_len):
+                    input_start = cum_seq_lens[i].item()
+                    input_pos = input_start + pos
+
+                    # Compare input and output for valid positions
+                    input_data = q_input[input_pos]
+                    output_data = padded_q[i, pos]
+
+                    torch.testing.assert_close(
+                        input_data, output_data, rtol=1e-5, atol=1e-6
+                    )
+
+                # Check that invalid positions are zero
+                for pos in range(seq_len, max_seq_len):
+                    output_data = padded_q[i, pos]
+                    self.assertTrue(
+                        torch.allclose(output_data, torch.zeros_like(output_data)),
+                        f"Position {pos} in batch {i} should be zero",
+                    )
+
+        # Test 2: unpad_draft_extend_output_kernel basic functionality
+        with self.subTest(test="unpad_kernel_basic"):
+            batch_size = 4
+            token_per_batch = 8
+            tp_q_head_num = 16
+            v_head_dim = 64
+
+            raw_out, output, accept_lengths, cum_accept_lengths = (
+                _create_test_output_data(
+                    self, batch_size, token_per_batch, tp_q_head_num, v_head_dim
+                )
+            )
+
+            # Launch kernel
+            BLOCK_SIZE = 64
+            grid = (batch_size * token_per_batch,)
+
+            unpad_draft_extend_output_kernel[grid](
+                raw_out_ptr=raw_out,
+                output_ptr=output,
+                accept_length_ptr=accept_lengths,
+                cumsum_ptr=cum_accept_lengths,
+                batch_size=batch_size,
+                token_per_batch=token_per_batch,
+                tp_q_head_num=tp_q_head_num,
+                v_head_dim=v_head_dim,
+                BLOCK_SIZE=BLOCK_SIZE,
+            )
+
+            # Verify the unpadding worked correctly
+            for i in range(batch_size):
+                accept_len = accept_lengths[i].item()
+                output_start = cum_accept_lengths[i].item()
+
+                # Check that valid positions are copied correctly
+                for pos in range(accept_len):
+                    input_data = raw_out[i, pos]
+                    output_data = output[output_start + pos]
+
+                    torch.testing.assert_close(
+                        input_data, output_data, rtol=1e-5, atol=1e-6
+                    )
+
 
 if __name__ == "__main__":
     unittest.main()

sglang/test/few_shot_gsm8k_engine.py

@@ -1,16 +1,14 @@
 import argparse
 import ast
 import asyncio
-import json
 import re
 import time
+from typing import Optional
 
 import numpy as np
 
 import sglang as sgl
-from sglang.lang.api import set_default_backend
-from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
-from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl
+from sglang.utils import download_and_cache_file, read_jsonl
 
 INVALID = -9999999