⚡️ Speed up method HybridLinearKVPool._transfer_full_attention_id by 18%

[codeflash-ai]

📄 18% (0.18x) speedup for HybridLinearKVPool._transfer_full_attention_id in python/sglang/srt/mem_cache/memory_pool.py

⏱️ Runtime : 11.0 microseconds → 9.26 microseconds (best of 10 runs)

📝 Explanation and details

The optimization replaces a conditional check-then-lookup pattern with a direct try/except approach for dictionary access. The original code uses if layer_id not in self.full_attention_layer_id_mapping: followed by a separate dictionary lookup, which results in two dictionary operations - one for the membership test and another for the actual value retrieval.

The optimized version uses try/except KeyError which performs only one dictionary lookup in the success case. In Python, dictionary __getitem__ is highly optimized and faster than separate __contains__ + __getitem__ calls.

Key changes:

• Eliminated double dictionary lookup by using try/except pattern
• In the error path, converts dict.keys() to a list directly for cleaner string formatting (dict_keys views are slower to stringify)

Why it's faster:

• Success path: One dict lookup vs two operations (18% speedup comes primarily from this)
• Exception path: Slightly optimized error message formatting
• Python's try/except is very efficient when exceptions are infrequent

Test case performance:
The optimization particularly benefits scenarios where _transfer_full_attention_id is called frequently with valid layer_ids (the common case), as seen in the test cases where invalid lookups are the minority. The single lookup approach provides consistent performance gains across all valid access patterns.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

from typing import List

# imports
import pytest
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool


# Dummy base classes and dependencies for isolated testing
class KVCache:
    pass

class MambaPool:
    pass

class MHATokenToKVPool:
    def __init__(self, size, page_size, dtype, head_num, head_dim, layer_num, device, enable_memory_saver):
        self.size = size
        self.page_size = page_size
        self.dtype = dtype
        self.head_num = head_num
        self.head_dim = head_dim
        self.layer_num = layer_num
        self.device = device
        self.enable_memory_saver = enable_memory_saver

    def get_kv_size_bytes(self):
        # Return dummy values for testing
        return (self.size * self.head_num * self.head_dim, self.size * self.head_num * self.head_dim)

GB = 1024 * 1024 * 1024
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool

# unit tests

# ----------- Basic Test Cases -----------




def test_empty_full_attention_layer_ids():
    """Test with empty full_attention_layer_ids list."""
    pool = HybridLinearKVPool(
        size=1,
        dtype='float32',
        page_size=1,
        head_num=1,
        head_dim=1,
        full_attention_layer_ids=[],
        enable_kvcache_transpose=False,
        device='cpu',
        mamba_pool=MambaPool()
    )
    # Any layer_id should raise ValueError
    with pytest.raises(ValueError):
        pool._transfer_full_attention_id(0)
    with pytest.raises(ValueError):
        pool._transfer_full_attention_id(-1)
    with pytest.raises(ValueError):
        pool._transfer_full_attention_id(100)









#------------------------------------------------
from typing import List

# imports
import pytest  # used for our unit tests
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool


class DummyTokenToKVPool:
    def __init__(self, **kwargs):
        pass
    def get_kv_size_bytes(self):
        return (1024, 1024)

# Simulate is_npu() and dependent classes for testing
def is_npu():
    return False

MHATokenToKVPool = DummyTokenToKVPool
AscendTokenToKVPool = DummyTokenToKVPool

GB = 1024 * 1024 * 1024
_is_npu = is_npu()

class KVCache:
    pass

class MambaPool:
    pass
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool

# unit tests

# Basic Test Cases



def test_edge_empty_full_attention_layer_ids():
    # No full attention layers
    pool = HybridLinearKVPool(
        size=10, dtype="float32", page_size=1, head_num=1, head_dim=1,
        full_attention_layer_ids=[], enable_kvcache_transpose=False,
        device="cpu", mamba_pool=MambaPool()
    )
    # Any layer_id should raise
    with pytest.raises(ValueError):
        pool._transfer_full_attention_id(0)
    with pytest.raises(ValueError):
        pool._transfer_full_attention_id(-1)
    with pytest.raises(ValueError):
        pool._transfer_full_attention_id(999)

To edit these changes git checkout codeflash/optimize-HybridLinearKVPool._transfer_full_attention_id-mh2mcnwn and push.