sglang 0.5.2rc2__py3-none-any.whl → 0.5.3.post1__py3-none-any.whl

sglang/srt/weight_sync/utils.py

@@ -6,7 +6,7 @@ from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import DTensor
 
 from sglang.srt.entrypoints.engine import Engine
-from sglang.srt.managers.tokenizer_manager import UpdateWeightsFromTensorReqInput
+from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
 from sglang.srt.model_executor.model_runner import LocalSerializedTensor
 from sglang.srt.utils import MultiprocessingSerializer
 
@@ -33,7 +33,7 @@ async def update_weights(
     """
     infer_tp_size = device_mesh[device_mesh_key].mesh.size()[0]
     infer_tp_rank = device_mesh[device_mesh_key].get_local_rank()
-    from sglang.srt.patch_torch import monkey_patch_torch_reductions
+    from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
 
     monkey_patch_torch_reductions()
 

sglang/test/attention/test_trtllm_mla_backend.py

@@ -41,6 +41,10 @@ DEFAULT_CONFIG = {
     "v_head_dim": 512,
     "num_kv_heads": 1,
     "layer_id": 0,
+    "tp_q_head_num": 128,
+    "tp_k_head_num": 128,
+    "prefill_head_dim": 192,
+    "prefill_v_head_dim": 128,
 }
 
 ROPE_BASE = 10000
@@ -92,7 +96,7 @@ TEST_CASES = {
             "description": "Medium-scale batch",
         },
     ],
-    "decode_output_match": [
+    "output_match": [
         {
             "name": "single_fp16",
             "batch_size": 1,
@@ -322,7 +326,7 @@ class TestTRTLLMMLA(CustomTestCase):
         config.update(test_case)
         return config
 
-    def _create_model_components(self, config):
+    def _create_model_components(self, config, is_prefill=False):
         """Create model runners, backends, and layer for testing."""
         # Create model runners
         model_runner_trtllm = MockModelRunner(config)
@@ -332,14 +336,23 @@ class TestTRTLLMMLA(CustomTestCase):
         trtllm_backend = TRTLLMMLABackend(model_runner_trtllm)
         reference_backend = FlashInferMLAAttnBackend(model_runner_reference)
 
+        head_dim = (
+            config["kv_lora_rank"] + config["qk_rope_head_dim"]
+            if not is_prefill
+            else config["prefill_head_dim"]
+        )
+        v_head_dim = (
+            config["v_head_dim"] if not is_prefill else config["prefill_v_head_dim"]
+        )
+
         # Create RadixAttention layer
         layer = RadixAttention(
             num_heads=config["num_attention_heads"],
-            head_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"],
+            head_dim=head_dim,
             scaling=model_runner_trtllm.model_config.scaling,
             num_kv_heads=config["num_kv_heads"],
             layer_id=config["layer_id"],
-            v_head_dim=config["v_head_dim"],
+            v_head_dim=v_head_dim,
             prefix="attn_mqa",
         )
 
@@ -524,7 +537,7 @@ class TestTRTLLMMLA(CustomTestCase):
         """Test that TRTLLM and FlashInfer MLA backends produce matching outputs."""
         print(f"\nRunning decode output matching tests...")
 
-        for test_case in TEST_CASES["decode_output_match"]:
+        for test_case in TEST_CASES["output_match"]:
             with self.subTest(test_case=test_case["name"]):
                 print(f"  Testing {test_case['name']}: {test_case['description']}")
 
@@ -1099,6 +1112,157 @@ class TestTRTLLMMLA(CustomTestCase):
         self.assertIsNotNone(metadata_3.block_kv_indices)
         self.assertEqual(metadata_3.block_kv_indices.shape[0], config["batch_size"])
 
+    def test_prefill_output_match_self_attention(self):
+        """Test prefill (forward) behavior of TRTLLM MLA backend vs reference."""
+        print(f"\nRunning prefill output tests...")
+
+        for test_case in TEST_CASES["output_match"][:2]:  # Just a subset for speed
+            with self.subTest(test_case=test_case["name"]):
+                print(
+                    f"Prefill Testing {test_case['name']}: {test_case['description']}"
+                )
+
+                config = self._merge_config(test_case)
+                batch_size = config["batch_size"]
+                max_seq_len = config["max_seq_len"]
+
+                # Create components
+                (
+                    model_runner_trtllm,
+                    model_runner_reference,
+                    trtllm_backend,
+                    reference_backend,
+                    layer,
+                ) = self._create_model_components(config, is_prefill=True)
+
+                # Prefill uses full sequences
+                seq_lens = torch.full(
+                    (batch_size,), max_seq_len, device=config["device"]
+                )
+
+                def _create_forward_batch_prefill(
+                    batch_size,
+                    seq_lens,
+                    extend_prefix_lens,
+                    backend,
+                    model_runner,
+                    config,
+                ):
+                    """Create a forward batch for the given backend."""
+
+                    fb = ForwardBatch(
+                        batch_size=batch_size,
+                        input_ids=torch.randint(
+                            0, 100, (batch_size, 1), device=config["device"]
+                        ),
+                        out_cache_loc=torch.arange(batch_size, device=config["device"]),
+                        seq_lens_sum=int(seq_lens.sum().item()),
+                        extend_prefix_lens=extend_prefix_lens,
+                        extend_prefix_lens_cpu=extend_prefix_lens.cpu().int().tolist(),
+                        extend_seq_lens_cpu=(seq_lens - extend_prefix_lens)
+                        .cpu()
+                        .int()
+                        .tolist(),
+                        forward_mode=ForwardMode.EXTEND,
+                        req_pool_indices=torch.arange(
+                            batch_size, device=config["device"]
+                        ),
+                        seq_lens=seq_lens,
+                        seq_lens_cpu=seq_lens.cpu(),
+                        attn_attend_prefix_cache=False,
+                        mha_return_lse=False,
+                        attn_backend=backend,
+                    )
+                    fb.req_to_token_pool = model_runner.req_to_token_pool
+                    fb.token_to_kv_pool = model_runner.token_to_kv_pool
+
+                    # Add position information for RoPE
+                    fb.positions = torch.arange(batch_size, device=config["device"])
+
+                    return fb
+
+                # Create forward batches
+                fb_trtllm = _create_forward_batch_prefill(
+                    batch_size,
+                    seq_lens.clone(),
+                    torch.zeros(batch_size, device=config["device"], dtype=torch.int32),
+                    trtllm_backend,
+                    model_runner_trtllm,
+                    config,
+                )
+                fb_reference = _create_forward_batch_prefill(
+                    batch_size,
+                    seq_lens.clone(),
+                    torch.zeros(batch_size, device=config["device"], dtype=torch.int32),
+                    reference_backend,
+                    model_runner_reference,
+                    config,
+                )
+
+                # Initialize metadata for both backends
+                trtllm_backend.init_forward_metadata(fb_trtllm)
+                reference_backend.init_forward_metadata(fb_reference)
+
+                # Create Q, K, V tensors for prefill
+                torch.manual_seed(config["seed_qkv"])
+
+                def _create_qkv_tensors_prefill(
+                    batch_size, seq_len, config, dtype_override=None
+                ):
+                    """Create Q, K, V tensors for prefill, using config for head_num and head_dim."""
+                    device = config["device"]
+                    dtype = dtype_override or config["dtype"]
+
+                    total_tokens = batch_size * seq_len
+
+                    tp_q_head_num = config["tp_q_head_num"]
+                    tp_k_head_num = config["tp_k_head_num"]
+                    head_dim = config["prefill_head_dim"]
+                    v_head_dim = config["prefill_v_head_dim"]
+
+                    q = torch.randn(
+                        (total_tokens, tp_q_head_num * head_dim),
+                        dtype=dtype,
+                        device=device,
+                    )
+                    k = torch.randn(
+                        (total_tokens, tp_k_head_num * head_dim),
+                        dtype=dtype,
+                        device=device,
+                    )
+                    v = torch.randn(
+                        (total_tokens, tp_k_head_num * v_head_dim),
+                        dtype=dtype,
+                        device=device,
+                    )
+
+                    # Reshape as requested
+                    q = q.view(-1, tp_q_head_num, head_dim)
+                    k = k.view(-1, tp_k_head_num, head_dim)
+                    v = v.view(-1, tp_k_head_num, v_head_dim)
+
+                    return q, k, v
+
+                q, k, v = _create_qkv_tensors_prefill(batch_size, max_seq_len, config)
+                # Run prefill on both backends
+                out_trtllm = trtllm_backend.forward_extend(
+                    q, k, v, layer, fb_trtllm, False
+                ).view(-1, layer.tp_q_head_num * layer.v_head_dim)
+                out_reference = reference_backend.forward_extend(
+                    q, k, v, layer, fb_reference, False
+                )
+
+                tolerance = config.get("tolerance", 1e-2)
+                comparison_passed = compare_outputs(
+                    out_trtllm, out_reference, tolerance=tolerance
+                )
+                self.assertTrue(
+                    comparison_passed,
+                    f"TRTLLM and Reference prefill outputs differ beyond tolerance. "
+                    f"Config: {test_case['name']}, "
+                    f"Max diff: {(out_trtllm - out_reference).abs().max().item()}",
+                )
+
 
 if __name__ == "__main__":
     unittest.main()

sglang/test/get_logits_ut.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+
+
+class DummyModel(nn.Module):
+    def __init__(self, d_in=2048, n_heads=128, softmax_scale=0.5):
+        super().__init__()
+        self.weights_proj = nn.Linear(d_in, 1024)
+        self.n_heads = n_heads
+        self.softmax_scale = softmax_scale
+
+    def _get_logits_head_gate_orig(self, x: torch.Tensor, q_scale: torch.Tensor):
+        weights = self.weights_proj(x)
+        weights = weights * self.n_heads**-0.5
+        q_scale = q_scale.unsqueeze(1)  # (B,1,1)
+        weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
+        return weights
+
+    def _get_logits_head_gate_opt(self, x: torch.Tensor, q_scale: torch.Tensor):
+        weights = self.weights_proj(x)
+        q_scale = q_scale.unsqueeze(1)  # (B,1,1)
+        scale_const = self.n_heads**-0.5 * q_scale * self.softmax_scale  # (B,1,1)
+        weights = weights.unsqueeze(-1) * scale_const  # (B,1024,1)
+        return weights
+
+
+def main():
+    torch.manual_seed(0)
+    model = DummyModel(d_in=2048, n_heads=128, softmax_scale=0.5)
+    x = torch.randn(128, 2048)  # batch=128, d_in=2048
+    q_scale = torch.randn(128, 1)
+
+    import time
+
+    start = time.time()
+    for _ in range(1000):
+        out_orig = model._get_logits_head_gate_orig(x, q_scale)
+    print("Original version time:", time.time() - start)
+
+    start = time.time()
+    for _ in range(1000):
+        out_opt = model._get_logits_head_gate_opt(x, q_scale)
+    print("Optimized version time:", time.time() - start)
+
+    print("Difference:", (out_orig - out_opt).abs().max().item())
+    assert torch.allclose(out_orig, out_opt), "Mismatch between original and optimized"
+
+
+if __name__ == "__main__":
+    main()
+
+
+"""
+Original version time: 0.49235057830810547
+Optimized version time: 0.4087331295013428
+Difference: 1.4901161193847656e-08
+"""

sglang/test/longbench_v2/__init__.py

@@ -0,0 +1 @@
+"""LongBench-v2 auxiliary utilities and validation scripts."""

sglang/test/longbench_v2/test_longbench_v2_eval.py

@@ -0,0 +1,238 @@
+"""
+Test cases for LongBench-v2 evaluation utility.
+"""
+
+import json
+import os
+import tempfile
+
+from sglang.test.simple_eval_longbench_v2 import (
+    LongBenchV2Eval,
+    extract_longbench_v2_answer,
+    format_longbench_v2_question,
+)
+
+
+def test_format_longbench_v2_question():
+    """Test the official LongBench-v2 question formatting."""
+    sample_row = {
+        "context": "This is a sample context about environmental issues.",
+        "question": "What is the main theme?",
+        "A": "Technology",
+        "B": "Environment",
+        "C": "Economics",
+        "D": "Politics",
+        "answer": "B",
+    }
+
+    formatted = format_longbench_v2_question(sample_row)
+
+    # Verify official template structure
+    assert "This is a sample context about environmental issues." in formatted
+    assert (
+        "What is the correct answer to this question: What is the main theme?"
+        in formatted
+    )
+    assert "(A) Technology" in formatted
+    assert "(B) Environment" in formatted
+    assert "(C) Economics" in formatted
+    assert "(D) Politics" in formatted
+    assert "The correct answer is" in formatted
+    print("✓ Question formatting works correctly")
+
+
+def test_extract_longbench_v2_answer():
+    """Test the official LongBench-v2 answer extraction."""
+
+    # Test official format: "The correct answer is (A)"
+    response1 = "After analyzing the context, The correct answer is (B)."
+    assert extract_longbench_v2_answer(response1) == "B"
+
+    # Test alternative format: "The correct answer is A"
+    response2 = "Based on the evidence, The correct answer is C."
+    assert extract_longbench_v2_answer(response2) == "C"
+
+    # Test with asterisks
+    response3 = "*The correct answer is (D)*"
+    assert extract_longbench_v2_answer(response3) == "D"
+
+    # Test fallback to standard pattern
+    response4 = "I think the answer is A."
+    assert extract_longbench_v2_answer(response4) == "A"
+
+    # Test no answer
+    response5 = "I'm not sure about this."
+    assert extract_longbench_v2_answer(response5) is None
+
+    print("✓ Answer extraction works correctly")
+
+
+def test_longbench_v2_eval_initialization():
+    """Test LongBench-v2 evaluation class initialization."""
+
+    # Create a temporary JSON file with sample data
+    sample_data = [
+        {
+            "_id": "test_001",
+            "domain": "single_document_qa",
+            "question": "What is X?",
+            "choice_A": "Option A1",
+            "choice_B": "Option B1",
+            "choice_C": "Option C1",
+            "choice_D": "Option D1",
+            "answer": "A",
+            "context": "Context 1",
+        },
+        {
+            "_id": "test_002",
+            "domain": "multi_document_qa",
+            "question": "What is Y?",
+            "A": "Option A2",
+            "B": "Option B2",
+            "C": "Option C2",
+            "D": "Option D2",
+            "answer": "B",
+            "context": "Context 2",
+        },
+    ]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(sample_data, f)
+        temp_file = f.name
+
+    try:
+        # Test initialization with new data_source parameter
+        eval_instance = LongBenchV2Eval(data_source=temp_file, num_examples=1)
+        assert len(eval_instance.examples) == 1
+        first_example = eval_instance.examples[0]
+        assert first_example.get("category") in {
+            "single_document_qa",
+            "multi_document_qa",
+        }
+        assert first_example.get("A") in {"Option A1", "Option A2"}
+        print("✓ Evaluation class initialization works correctly")
+
+    finally:
+        os.unlink(temp_file)
+
+
+def test_category_filtering():
+    """Ensure category filtering keeps only requested domains."""
+
+    sample_data = [
+        {
+            "_id": "test_001",
+            "domain": "single_document_qa",
+            "question": "What is X?",
+            "choice_A": "Option A1",
+            "choice_B": "Option B1",
+            "choice_C": "Option C1",
+            "choice_D": "Option D1",
+            "answer": "A",
+            "context": "Context 1",
+        },
+        {
+            "_id": "test_002",
+            "domain": "multi_document_qa",
+            "question": "What is Y?",
+            "choice_A": "Option A2",
+            "choice_B": "Option B2",
+            "choice_C": "Option C2",
+            "choice_D": "Option D2",
+            "answer": "B",
+            "context": "Context 2",
+        },
+    ]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(sample_data, f)
+        temp_file = f.name
+
+    try:
+        eval_instance = LongBenchV2Eval(
+            data_source=temp_file,
+            categories=["multi_document_qa"],
+        )
+        assert len(eval_instance.examples) == 1
+        assert eval_instance.examples[0]["category"] == "multi_document_qa"
+        print("✓ Category filtering works correctly")
+    finally:
+        os.unlink(temp_file)
+
+
+def test_difficulty_metrics():
+    """Validate that difficulty-specific metrics are recorded."""
+
+    sample_data = [
+        {
+            "_id": "easy_001",
+            "domain": "single_document_qa",
+            "difficulty": "easy",
+            "question": "Easy question?",
+            "choice_A": "Correct",
+            "choice_B": "Wrong",
+            "choice_C": "Wrong",
+            "choice_D": "Wrong",
+            "answer": "A",
+            "context": "Easy context",
+        },
+        {
+            "_id": "hard_001",
+            "domain": "single_document_qa",
+            "difficulty": "hard",
+            "question": "Hard question?",
+            "choice_A": "Wrong",
+            "choice_B": "Correct",
+            "choice_C": "Wrong",
+            "choice_D": "Wrong",
+            "answer": "B",
+            "context": "Hard context",
+        },
+    ]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(sample_data, f)
+        temp_file = f.name
+
+    class FixedSampler:  # noqa: D401 - simple helper
+        """Mock sampler returning the correct answer based on question text."""
+
+        def _pack_message(self, content: str, role: str):
+            return {"content": content, "role": role}
+
+        def __call__(self, messages):
+            prompt = messages[0]["content"]
+            if "Easy question" in prompt:
+                return "The correct answer is (A)"
+            return "The correct answer is (B)"
+
+    try:
+        eval_instance = LongBenchV2Eval(data_source=temp_file, num_threads=1)
+        result = eval_instance(FixedSampler())
+
+        assert result.metrics.get("difficulty_easy") == 1.0
+        assert result.metrics.get("difficulty_hard") == 1.0
+        print("✓ Difficulty metrics recorded correctly")
+    finally:
+        os.unlink(temp_file)
+
+
+def main():
+    """Run all tests."""
+    print("Testing simplified LongBench-v2 evaluation utility...\n")
+
+    test_format_longbench_v2_question()
+    test_extract_longbench_v2_answer()
+    test_longbench_v2_eval_initialization()
+    test_category_filtering()
+    test_difficulty_metrics()
+
+    print("\n" + "=" * 50)
+    print("✅ ALL TESTS PASSED!")
+    print("The simplified implementation follows SGLang patterns")
+    print("while maintaining LongBench-v2 compatibility.")
+    print("=" * 50)
+
+
+if __name__ == "__main__":
+    main()