⚡️ Speed up function _topk_ids_logical_to_physical_dynamic by 12%

[codeflash-ai]

📄 12% (0.12x) speedup for _topk_ids_logical_to_physical_dynamic in python/sglang/srt/eplb/expert_location_dispatch.py

⏱️ Runtime : 1.71 milliseconds → 1.52 milliseconds (best of 247 runs)

📝 Explanation and details

The optimized code achieves a 12% speedup through several key tensor operation optimizations:

Key Optimizations:

1. Smart reshape handling: Replaces unconditional .flatten() with conditional .reshape(-1) only when needed (multi-dimensional tensors), avoiding unnecessary memory copies for already 1D tensors.

2. Efficient random generation: Uses torch.empty(...).random_(65536) instead of torch.randint(0, 65536, ...), which generates random numbers in-place rather than creating intermediate tensors, reducing memory allocation overhead.

3. Variable caching: Stores the flattened indices in topk_idx to avoid recomputation during indexing operations.

4. Conditional view operation: Only calls .view() to restore original shape when the shape actually changed, eliminating unnecessary tensor operations.

Performance Impact:
The function is called from topk_ids_logical_to_physical() for "dynamic" and "fake" expert dispatch algorithms, suggesting it's in a hot path for expert routing in distributed inference scenarios. The optimizations are particularly effective for:

• Large batches: Test results show 7-22% speedups for scenarios with many logical IDs (1000+ elements)
• 1D tensors: Up to 28% improvement when input is already flat
• Multiple physical mappings: 19-21% gains when experts have multiple valid physical locations

The optimizations preserve exact functionality while reducing memory allocations and tensor operations, making expert routing more efficient in distributed model serving workloads where this function may be called frequently during request processing.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 49 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
import torch
from sglang.srt.eplb.expert_location_dispatch import \
    _topk_ids_logical_to_physical_dynamic


# Helper class to mock ExpertLocationDispatchInfo
class ExpertLocationDispatchInfo:
    def __init__(self, partial_logical_to_all_physical_map, partial_logical_to_all_physical_map_num_valid):
        self.partial_logical_to_all_physical_map = partial_logical_to_all_physical_map
        self.partial_logical_to_all_physical_map_num_valid = partial_logical_to_all_physical_map_num_valid
from sglang.srt.eplb.expert_location_dispatch import \
    _topk_ids_logical_to_physical_dynamic

# ========== UNIT TESTS ==========

# ----------- BASIC TEST CASES -----------
def test_basic_single_id_single_physical():
    # One logical ID, one valid physical mapping
    logical_to_physical_map = torch.tensor([[42]], dtype=torch.int64)
    num_valid = torch.tensor([1], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 50.1μs -> 43.1μs (16.2% faster)

def test_basic_multiple_ids_single_physical_each():
    # Multiple logical IDs, each maps to one physical
    logical_to_physical_map = torch.tensor([[1], [2], [3]], dtype=torch.int64)
    num_valid = torch.tensor([1, 1, 1], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0, 1, 2], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 43.1μs -> 35.4μs (21.8% faster)

def test_basic_multiple_ids_multiple_physical():
    # Logical IDs map to multiple physical, random selection
    logical_to_physical_map = torch.tensor([[10, 11], [20, 21]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0, 1], dtype=torch.int64)
    # Each output must be in the corresponding row
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 44.5μs -> 37.1μs (19.8% faster)

def test_basic_batch_shape_preserved():
    # Input is 2D, shape must be preserved
    logical_to_physical_map = torch.tensor([[5, 6], [7, 8]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 43.9μs -> 44.7μs (1.86% slower)
    # Each element must be from the corresponding logical id's valid physicals
    for i in range(2):
        for j in range(2):
            logical_id = topk_ids[i, j].item()

# ----------- EDGE TEST CASES -----------
def test_edge_empty_input():
    # Empty input tensor
    logical_to_physical_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.empty((0,), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 36.7μs -> 30.9μs (19.1% faster)

def test_edge_single_valid_multiple_invalid():
    # Logical IDs with some zero valid mappings
    logical_to_physical_map = torch.tensor([[100, 101], [200, 201]], dtype=torch.int64)
    num_valid = torch.tensor([2, 0], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 40.0μs -> 33.7μs (18.8% faster)
    # Now test with invalid logical id (should raise error due to zero valid)
    topk_ids_invalid = torch.tensor([1], dtype=torch.int64)
    with pytest.raises(RuntimeError):
        _topk_ids_logical_to_physical_dynamic(topk_ids_invalid, info) # 59.9μs -> 58.2μs (2.99% faster)

def test_edge_maximum_valid_index():
    # Logical ID with maximum valid index (test modulus logic)
    logical_to_physical_map = torch.tensor([[1, 2, 3]], dtype=torch.int64)
    num_valid = torch.tensor([3], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0], dtype=torch.int64)
    # Try multiple times to ensure all possible outputs
    outputs = set()
    for _ in range(20):
        codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 252μs -> 208μs (21.2% faster)
        outputs.add(result[0].item())

def test_edge_non_contiguous_input():
    # Non-contiguous input tensor
    logical_to_physical_map = torch.tensor([[10, 11], [20, 21]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    full = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
    topk_ids = full[::2]  # [0, 0]
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 41.2μs -> 37.1μs (11.1% faster)
    for val in result:
        pass

def test_edge_large_id_values():
    # Logical IDs with large values
    logical_to_physical_map = torch.cat([
        torch.arange(100000, 100002).unsqueeze(0),
        torch.arange(200000, 200002).unsqueeze(0)
    ], dim=0)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0, 1], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 36.5μs -> 30.8μs (18.6% faster)

def test_edge_device_cpu_and_cuda():
    # Test both CPU and CUDA (if available)
    logical_to_physical_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.tensor([0, 1], dtype=torch.int64)
    # CPU
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result_cpu = codeflash_output # 41.2μs -> 32.4μs (27.4% faster)
    # CUDA (if available)
    if torch.cuda.is_available():
        logical_to_physical_map_cuda = logical_to_physical_map.cuda()
        num_valid_cuda = num_valid.cuda()
        info_cuda = ExpertLocationDispatchInfo(logical_to_physical_map_cuda, num_valid_cuda)
        topk_ids_cuda = topk_ids.cuda()
        codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids_cuda, info_cuda); result_cuda = codeflash_output

# ----------- LARGE SCALE TEST CASES -----------
def test_large_scale_many_logical_ids():
    # 1000 logical IDs, each maps to 3 physical IDs
    num_logical = 1000
    num_physical = 3
    logical_to_physical_map = torch.arange(num_logical * num_physical, dtype=torch.int64).view(num_logical, num_physical)
    num_valid = torch.full((num_logical,), num_physical, dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.arange(num_logical, dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 59.8μs -> 55.5μs (7.82% faster)
    # Each result must be in the corresponding row
    for i in range(num_logical):
        pass

def test_large_scale_many_topk_ids():
    # 1000 topk_ids, all map to the same logical id
    logical_to_physical_map = torch.tensor([[7, 8, 9]], dtype=torch.int64)
    num_valid = torch.tensor([3], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.zeros(1000, dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 56.5μs -> 50.3μs (12.4% faster)
    for val in result:
        pass

def test_large_scale_2d_shape():
    # Large 2D shape, shape must be preserved
    logical_to_physical_map = torch.tensor([[4, 5], [6, 7]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.randint(0, 2, (50, 20), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 65.7μs -> 68.7μs (4.48% slower)
    for i in range(50):
        for j in range(20):
            logical_id = topk_ids[i, j].item()

def test_large_scale_randomized():
    # Randomized logical_to_physical_map and topk_ids
    torch.manual_seed(1234)
    num_logical = 100
    num_physical = 5
    logical_to_physical_map = torch.randint(1000, 2000, (num_logical, num_physical), dtype=torch.int64)
    num_valid = torch.randint(1, num_physical + 1, (num_logical,), dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.randint(0, num_logical, (200,), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 47.2μs -> 43.8μs (7.72% faster)
    for i in range(200):
        logical_id = topk_ids[i].item()
        valid_count = num_valid[logical_id].item()
        valid_physical = logical_to_physical_map[logical_id, :valid_count].tolist()

def test_large_scale_memory_limit():
    # Ensure memory usage is under 100MB (approximate)
    num_logical = 500
    num_physical = 10
    logical_to_physical_map = torch.randint(0, 100000, (num_logical, num_physical), dtype=torch.int64)
    num_valid = torch.randint(1, num_physical + 1, (num_logical,), dtype=torch.int64)
    info = ExpertLocationDispatchInfo(logical_to_physical_map, num_valid)
    topk_ids = torch.randint(0, num_logical, (1000,), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 56.9μs -> 53.7μs (6.05% faster)
    # Check that each output is valid
    for i in range(1000):
        logical_id = topk_ids[i].item()
        valid_count = num_valid[logical_id].item()
        valid_physical = logical_to_physical_map[logical_id, :valid_count].tolist()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest
import torch
from sglang.srt.eplb.expert_location_dispatch import \
    _topk_ids_logical_to_physical_dynamic


# Dummy class to simulate ExpertLocationDispatchInfo for the tests
class ExpertLocationDispatchInfo:
    def __init__(self, partial_logical_to_all_physical_map, partial_logical_to_all_physical_map_num_valid):
        self.partial_logical_to_all_physical_map = partial_logical_to_all_physical_map
        self.partial_logical_to_all_physical_map_num_valid = partial_logical_to_all_physical_map_num_valid
from sglang.srt.eplb.expert_location_dispatch import \
    _topk_ids_logical_to_physical_dynamic

# Basic Test Cases

def test_basic_single_id():
    # Test with a single id
    partial_map = torch.tensor([[100, 101], [200, 201]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 50.1μs -> 42.9μs (16.7% faster)

def test_basic_multiple_ids():
    # Test with multiple ids
    partial_map = torch.tensor([[10, 11], [20, 21], [30, 31]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0, 1, 2], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 42.7μs -> 35.5μs (20.3% faster)

def test_basic_2d_shape():
    # Test with a 2D shape
    partial_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 43.0μs -> 43.6μs (1.43% slower)
    # Each entry is valid
    for i in range(2):
        for j in range(2):
            logical_id = topk_ids[i, j].item()

# Edge Test Cases

def test_edge_single_valid_mapping():
    # Only one valid mapping per logical id
    partial_map = torch.tensor([[42, 0], [99, 0]], dtype=torch.int64)
    num_valid = torch.tensor([1, 1], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0, 1], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 42.5μs -> 34.1μs (24.8% faster)

def test_edge_zero_valid_mapping_raises():
    # Zero valid mapping should raise an error due to modulo by zero
    partial_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([0, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0], dtype=torch.int64)
    with pytest.raises(RuntimeError):
        _topk_ids_logical_to_physical_dynamic(topk_ids, info) # 82.2μs -> 77.0μs (6.79% faster)

def test_edge_non_contiguous_ids():
    # Logical ids are not contiguous and have different num_valid
    partial_map = torch.tensor([[5, 6], [7, 8], [9, 10]], dtype=torch.int64)
    num_valid = torch.tensor([2, 1, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([1, 2, 0], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 44.6μs -> 39.9μs (11.8% faster)

def test_edge_device_cpu_and_cuda():
    # Test on both CPU and CUDA (if available)
    partial_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0, 1], dtype=torch.int64)
    # CPU
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result_cpu = codeflash_output # 41.2μs -> 33.9μs (21.7% faster)
    # CUDA (if available)
    if torch.cuda.is_available():
        partial_map_cuda = partial_map.cuda()
        num_valid_cuda = num_valid.cuda()
        info_cuda = ExpertLocationDispatchInfo(partial_map_cuda, num_valid_cuda)
        topk_ids_cuda = topk_ids.cuda()
        codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids_cuda, info_cuda); result_cuda = codeflash_output

def test_edge_large_partial_map_num_valid():
    # num_valid is not the same for all logical ids
    partial_map = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int64)
    num_valid = torch.tensor([1, 2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0, 1, 2], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 50.8μs -> 41.6μs (22.3% faster)

def test_edge_empty_topk_ids():
    # Empty topk_ids should return an empty tensor with same shape
    partial_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([], dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 38.8μs -> 30.1μs (28.8% faster)

def test_edge_topk_ids_out_of_bounds_raises():
    # topk_ids outside valid range should raise an error
    partial_map = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
    num_valid = torch.tensor([2, 2], dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.tensor([0, 2], dtype=torch.int64)  # 2 is out of bounds
    with pytest.raises(IndexError):
        _topk_ids_logical_to_physical_dynamic(topk_ids, info) # 74.4μs -> 64.5μs (15.4% faster)

# Large Scale Test Cases

def test_large_scale_many_ids():
    # Test with a large number of ids
    n_ids = 1000
    n_phys = 5
    partial_map = torch.arange(n_ids * n_phys, dtype=torch.int64).view(n_ids, n_phys)
    num_valid = torch.full((n_ids,), n_phys, dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.arange(n_ids, dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 59.7μs -> 56.5μs (5.67% faster)
    # Each result must be in the corresponding row of partial_map
    for i in range(n_ids):
        pass

def test_large_scale_random_ids_and_valid():
    # Test with random ids and random num_valid
    n_ids = 500
    n_phys = 4
    partial_map = torch.arange(n_ids * n_phys, dtype=torch.int64).view(n_ids, n_phys)
    # Random num_valid between 1 and n_phys
    num_valid = torch.randint(1, n_phys + 1, (n_ids,), dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.randint(0, n_ids, (1000,), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 52.4μs -> 49.8μs (5.26% faster)
    # Each result must be in the valid range for its logical id
    for i in range(topk_ids.shape[0]):
        logical_id = topk_ids[i].item()
        valid_phys = partial_map[logical_id, :num_valid[logical_id]].tolist()

def test_large_scale_2d_tensor():
    # Test with a large 2D tensor
    n_ids = 100
    n_phys = 8
    partial_map = torch.arange(n_ids * n_phys, dtype=torch.int64).view(n_ids, n_phys)
    num_valid = torch.full((n_ids,), n_phys, dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.randint(0, n_ids, (32, 32), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 56.4μs -> 58.8μs (4.05% slower)
    for i in range(topk_ids.shape[0]):
        for j in range(topk_ids.shape[1]):
            logical_id = topk_ids[i, j].item()

def test_large_scale_max_memory_constraint():
    # Ensure we do not exceed the memory constraint (100MB)
    # Each int64 is 8 bytes, so max 12,500,000 elements
    # We'll use 1000x10 elements, well within the limit
    n_ids = 1000
    n_phys = 10
    partial_map = torch.arange(n_ids * n_phys, dtype=torch.int64).view(n_ids, n_phys)
    num_valid = torch.full((n_ids,), n_phys, dtype=torch.int64)
    info = ExpertLocationDispatchInfo(partial_map, num_valid)
    topk_ids = torch.randint(0, n_ids, (1000,), dtype=torch.int64)
    codeflash_output = _topk_ids_logical_to_physical_dynamic(topk_ids, info); result = codeflash_output # 53.3μs -> 50.1μs (6.31% faster)
    for i in range(topk_ids.shape[0]):
        logical_id = topk_ids[i].item()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_topk_ids_logical_to_physical_dynamic-mhospliq and push.