본문으로 건너뛰기

[sglang] Apple Silicon LLM 성능 향상: 슬라이딩 윈도우 KV 캐싱 및 인-그래프 샘플링 도입

PR 링크: sgl-project/sglang#34166 상태: Merged | 변경: +3501 / -517

들어가며

최근 LLM(거대 언어 모델)의 발전은 놀라운 속도로 이루어지고 있으며, 특히 Apple Silicon과 같은 온디바이스 AI 환경에서의 성능 최적화는 사용자 경험에 직접적인 영향을 미칩니다. sglang 레포지토리의 이 PR은 Apple Silicon 환경에서 MLX 백엔드를 사용하는 LLM의 두 가지 주요 병목 현상을 해결합니다. 첫째, 슬라이딩 윈도우 어텐션(Sliding-Window Attention, SWA)을 사용하는 모델에서 불필요하게 전체 KV 히스토리를 저장하여 발생하는 메모리 낭비를 줄입니다. 둘째, 토큰 선택 방식이 항상 가장 확률이 높은 토큰만을 선택하는 그리디(greedy) 방식에 국한되었던 문제를 해결하고, temperature, top-p, top-k 등의 샘플링 기법을 MLX 백엔드에서 지원하도록 개선합니다. 이 두 가지 개선 사항은 특히 gpt-oss와 같은 대규모 모델을 Apple Silicon에서 효율적으로 서빙하는 데 중요한 역할을 합니다.

코드 분석

Part 1: Window-bounded SWA KV storage

이 PR의 핵심적인 메모리 최적화는 슬라이딩 윈도우 어텐션을 사용하는 모델의 KV 캐시 관리 방식에 있습니다. 기존에는 슬라이딩 윈도우를 사용하더라도 KV 캐시 자체는 전체 컨텍스트 길이를 모두 저장하여 메모리 사용량이 높았습니다. 이 PR에서는 WindowedAttentionKVCache라는 새로운 클래스를 도입하여, 실제로 어텐션 계산에 필요한 최근 window 크기만큼의 토큰만 KV 캐시에 저장하도록 변경했습니다. 이는 고정 크기의 버퍼를 사용하여 O(1)의 시간 복잡도로 토큰을 추가하고 압축할 수 있게 합니다.

주요 변경 사항:

  1. WindowedAttentionKVCache 도입: 기존의 ContiguousAttentionKVCache는 전체 히스토리를 저장했지만, WindowedAttentionKVCache는 슬라이딩 윈도우 크기(window)만큼만 저장합니다.

    --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
+++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
@@ -141,9 +152,130 @@
 def write_token(self, k: mx.array, v: mx.array) -> None:
     self.values[:, :, self.offset : end, :] = v
     self.offset = end
-
-    def get_kv(self) -> tuple[mx.array, mx.array]:
-        """Return valid K/V: (1, n_kv_heads, offset, head_dim)."""
-        return self.keys[:, :, : self.offset, :], self.values[:, :, : self.offset, :]
+
+    def get_kv(self, window: int | None = None) -> tuple[mx.array, mx.array]:
+        """Return valid K/V: (1, n_kv_heads, min(offset, window), head_dim).
+
+        ``window`` keeps only the trailing window a sliding-window layer can
+        attend to.  Slicing here rather than slicing the full history and then
+        slicing again costs one op instead of two per request per layer.
+        """
+        start = 0 if window is None else max(0, self.offset - window)
+        return (
+            self.keys[:, :, start : self.offset, :],
+            self.values[:, :, start : self.offset, :],
+        )
+
+    def reset(self) -> None:
+        """Reset for reuse, keeping allocated buffers."""
+        self.offset = 0
+
+
+class WindowedAttentionKVCache:
+    """Sliding-window attention KV buffer for one request and one layer.
+
+    Holds the trailing ``window`` tokens plus the in-flight chunk, in
+    temporal order, instead of the full sequence.  ``offset`` stays
+    absolute (RoPE positions, decode bookkeeping); the dropped prefix
+    s

  2. KV 풀(Pool) 최적화: 기존에는 모든 레이어의 KV 캐시를 위한 공유 풀(MlxAttentionKVPool)이 존재했습니다. 이 PR에서는 슬라이딩 윈도우 모델의 경우, 공유 풀을 사용하지 않고 각 레이어별로 독립적인 WindowedAttentionKVCache를 사용하도록 변경했습니다. 이는 공유 풀의 할당 자체를 건너뛰게 하여 상당한 메모리를 절약합니다.

    --- a/python/sglang/srt/hardware_backend/mlx/aot.py
+++ b/python/sglang/srt/hardware_backend/mlx/aot.py
@@ -219,7 +214,10 @@
     req_ids: list[str],
     req_pool_idx: dict[str, int],
     req_to_token_pool: Any | None,

  •    layer_caches: list[list[ContiguousAttentionKVCache]],

  •    # Only .offset is read (absolute on every cache kind) and the slot

  •    # lookup is layer-agnostic; the wrapper gates the fused pool scatter

  •    # to full-attention layers.

  •    layer_caches: list[list[Any]],
    ) -> MlxAOTKernelContext: """Build optional AOT context for one batched decode step.""" if not aot_kernels.rope.enabled or kv_pool is None: 
    또한, `init_cache_pools` 함수에서 슬라이딩 윈도우 레이어가 있는 모델의 경우 공유 풀 할당을 건너뛰도록 수정되었습니다. 이는 상당한 메모리(`~6.6 GiB`)를 절약할 수 있습니다.
  1. 마스크 생성 최적화: make_attention_mask 함수에서 슬라이딩 윈도우가 적용될 때, offset + N <= window_size인 경우 일반적인 인과 마스크와 동일한 결과를 생성하므로, 불필요한 마스크 생성을 건너뛰도록 개선했습니다. 이는 mx.fast.scaled_dot_product_attention의 융합된 인과 마스크 경로를 활용하게 하여 성능을 향상시킵니다. 
    --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
+++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
@@ -21,8 +21,19 @@

def make_attention_mask(N, offset, return_array=False, window_size=None): layers pass it, including for N == 1) or windowed models silently fall back to full attention. """

  • if window_size is not None:

Part 2: In-graph Sampling (--mlx-enable-sampling)

이 PR은 MLX 백엔드에서 다양한 샘플링 전략(temperature, top-p, top-k 등)을 지원하기 위해 --mlx-enable-sampling 플래그를 도입했습니다. 이 기능은 MLX의 지연 그래프(lazy graph) 내에서 직접 구현되어, 기존의 오버랩 스케줄러와 호환성을 유지하면서 토큰 샘플링을 가능하게 합니다.

주요 변경 사항:

  1. hardware_backend/mlx/sampling.py 신규 추가: 이 파일은 MLX 연산을 사용하여 토큰 샘플링 로직을 구현합니다. Gumbel-max 샘플링, seeded sampling, top-k/top-p 필터링 등을 지원합니다.

    --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
+++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
@@ -141,9 +152,130 @@
 def write_token(self, k: mx.array, v: mx.array) -> None:
     self.values[:, :, self.offset : end, :] = v
     self.offset = end
-
-    def get_kv(self) -> tuple[mx.array, mx.array]:
-        """Return valid K/V: (1, n_kv_heads, offset, head_dim)."""
-        return self.keys[:, :, : self.offset, :], self.values[:, :, : self.offset, :]
+
+    def get_kv(self, window: int | None = None) -> tuple[mx.array, mx.array]:
+        """Return valid K/V: (1, n_kv_heads, min(offset, window), head_dim).
+
+        ``window`` keeps only the trailing window a sliding-window layer can
+        attend to.  Slicing here rather than slicing the full history and then
+        slicing again costs one op instead of two per request per layer.
+        """
+        start = 0 if window is None else max(0, self.offset - window)
+        return (
+            self.keys[:, :, start : self.offset, :],
+            self.values[:, :, start : self.offset, :],
+        )
+
+    def reset(self) -> None:
+        """Reset for reuse, keeping allocated buffers."""
+        self.offset = 0
+
+
+class WindowedAttentionKVCache:
+    """Sliding-window attention KV buffer for one request and one layer.
+
+    Holds the trailing ``window`` tokens plus the in-flight chunk, in
+    temporal order, instead of the full sequence.  ``offset`` stays
+    absolute (RoPE positions, decode bookkeeping); the dropped prefix
+    s

    (Note: The diff above shows the KV cache part, but the sampling logic itself resides in hardware_backend/mlx/sampling.py which is a new file.)

  2. --mlx-enable-sampling 플래그 추가: 서버 실행 시 이 플래그를 활성화하면 샘플링 기능이 작동합니다. 기본값은 비활성화 상태입니다.

    --- a/docs/docs/advanced_features/server_arguments.mdx
+++ b/docs/docs/advanced_features/server_arguments.mdx
@@ -673,6 +673,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>None</code></td>
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
  </tr>

  •    <tr>

  •  <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--mlx-enable-sampling`</td>

  •  <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>MLX backend only: sample decode tokens (temperature / top-k / top-p / min-p) instead of greedy argmax. Sampling runs inside the lazy MLX graph, so it works with the overlap scheduler; first tokens from prefill/extend are sampled too. Greedy requests keep exact argmax behavior. Also enables on the MLX path: grammar vocab masks and custom logit processors (these break decode chaining per step; custom processors run on pure-decode steps only), logit_bias, output logprobs (sampled token / top-k / token_ids; prompt input logprobs are not computed), NaN sanitization (SGLANG_SANITIZE_NAN_LOGITS), and per-request sampling_seed under --enable-deterministic-inference (deterministic within MLX only). Penalties are not applied.</td>

  •  <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>

  •  <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
  • `--constrained-json-whitespace-pattern` Regular expression to match whitespace characters in JSON strings. 
```

  1. 샘플링 비용 절감: mx.argsort와 같은 비용이 많이 드는 연산을 [batch, K] 차원으로 줄이고, 불필요한 softmax 및 log 연산을 건너뛰는 등의 최적화를 통해 샘플링 오버헤드를 최소화했습니다.

    --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
+++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py
@@ -141,9 +152,130 @@
 def write_token(self, k: mx.array, v: mx.array) -> None:
     self.values[:, :, self.offset : end, :] = v
     self.offset = end
-
-    def get_kv(self) -> tuple[mx.array, mx.array]:
-        """Return valid K/V: (1, n_kv_heads, offset, head_dim)."""
-        return self.keys[:, :, : self.offset, :], self.values[:, :, : self.offset, :]
+
+    def get_kv(self, window: int | None = None) -> tuple[mx.array, mx.array]:
+        """Return valid K/V: (1, n_kv_heads, min(offset, window), head_dim).
+
+        ``window`` keeps only the trailing window a sliding-window layer can
+        attend to.  Slicing here rather than slicing the full history and then
+        slicing again costs one op instead of two per request per layer.
+        """
+        start = 0 if window is None else max(0, self.offset - window)
+        return (
+            self.keys[:, :, start : self.offset, :],
+            self.values[:, :, start : self.offset, :],
+        )
+
+    def reset(self) -> None:
+        """Reset for reuse, keeping allocated buffers."""
+        self.offset = 0
+
+
+class WindowedAttentionKVCache:
+    """Sliding-window attention KV buffer for one request and one layer.
+
+    Holds the trailing ``window`` tokens plus the in-flight chunk, in
+    temporal order, instead of the full sequence.  ``offset`` stays
+    absolute (RoPE positions, decode bookkeeping); the dropped prefix
+    s

    (Note: The diff above shows the KV cache part, but the sampling logic itself resides in hardware_backend/mlx/sampling.py which is a new file.)

  2. 그리디(Greedy) 단축 경로: 샘플링이 비활성화되거나 temperature=0 등 그리디 방식이 적용될 경우, 기존의 mx.argmax 연산을 그대로 사용하여 성능 저하 없이 정확도를 유지합니다.

기타 개선 사항

  • 디코드 패딩 최적화: 패딩 계산을 레이어별이 아닌 스텝별로 한 번만 수행하도록 변경하여 성능을 개선했습니다.
  • 마스크 생성 최적화: make_mask 함수에서 불필요한 마스크 생성을 건너뛰어 scaled_dot_product_attention의 융합된 인과 마스크 경로를 활용하도록 했습니다.
  • constrained/xgrammar_backend.py 수정: MLX 백엔드에서 pin_memory 및 CPU 기반 apply_vocab_mask 처리를 지원하도록 수정하여 CUDA 없는 환경에서의 크래시를 방지합니다.
  • logprob_result_processor.py 수정: next_token_top_logprobs_val 접근 시 top_logprobs_num > 0 조건을 확인하도록 하여 TypeError를 방지합니다.

왜 이게 좋은가?

메모리 절약

이 PR의 가장 큰 장점은 슬라이딩 윈도우 어텐션을 사용하는 모델의 메모리 사용량을 크게 줄였다는 점입니다. PR에서 제시된 gpt-oss-20b 모델의 측정 결과는 다음과 같습니다:

sequence full-history windowed saving
<= 4096 (preallocated) 192.0 MiB 102.0 MiB 46.9%
4096 prompt + 4000 decode 384.0 MiB 198.0 MiB 48.4%
8192 prompt + 192 decode 768.0 MiB 390.0 MiB 49.2%

이는 특히 긴 컨텍스트를 처리해야 하는 모델에서 Apple Silicon 장치의 제한된 메모리를 훨씬 효율적으로 사용할 수 있게 해줍니다. 또한, 공유 KV 풀 할당을 건너뛰면서 약 6.6 GiB의 메모리를 추가로 확보할 수 있습니다.

추론 속도 향상 및 기능 확장

  1. 샘플링 기능 도입: --mlx-enable-sampling 플래그를 통해 temperature, top-p, top-k 등의 샘플링 기법을 MLX 백엔드에서 지원하게 되면서, 더 다양하고 창의적인 텍스트 생성이 가능해졌습니다. PR에서는 이 기능이 기존 오버랩 스케줄러와 호환되면서도 성능 저하를 최소화했음을 보여줍니다. 샘플링 활성화 시 TPOT(Tokens Per Output Token)이 약 4.9% 증가하지만, 이는 MLX 그래프 내에서 처리되어 효율적입니다. 
    | arm                                             | tok/s  | TPOT ms | TTFT ms |
|-------------------------------------------------|--------|---------|---------|
| flag on, temperature 0.8 (sampling active)      | 810.4  | 4.311   | 71.4    |
| flag on, temperature 0 (greedy short-circuit)   | 823.3  | 4.108   | 71.8    |
| flag off (pre-PR baseline)                      | 817.8  | 4.196   | 70.9    |
  2. 마스크 생성 최적화: 슬라이딩 윈도우 모델에서 불필요한 마스크 생성을 건너뛰는 최적화는 약 2배의 속도 향상을 가져옵니다. 
    | case                 | banded  | cheap form | speedup |
|----------------------|---------|------------|---------|
| window 1024, N=1024  | 2.89 ms | 1.52 ms    | 2.03x   |
| window 4096, N=2048  | 10.90 ms| 5.38 ms    | 2.07x   |
| window 4096, N=4096  | 42.86 ms| 20.65 ms   | 2.10x   |
  3. 디코드 패딩 최적화: 디코드 패딩 계산을 스텝별로 한 번만 수행하도록 변경하여, ragged 배치에서 최대 11.5%의 성능 향상을 보였습니다. 
    | batch                | before  | after   | delta   |
|----------------------|---------|---------|---------|
| B=8 ragged 60-4096   | 39.47 ms| 34.92 ms| -11.5%  |
| B=8 ragged all > window | 38.59 ms| 35.74 ms| -7.4%   |

일반적인 교훈

  • 하드웨어 특화 최적화: Apple Silicon의 MLX 백엔드와 같이 특정 하드웨어의 특성을 깊이 이해하고 이를 활용하는 것이 성능 향상의 핵심입니다. 메모리 관리, 연산 그래프 최적화, 커널 융합 등 다양한 측면에서 최적화가 가능합니다.
  • 메모리 효율성: LLM의 KV 캐시는 메모리 사용량의 큰 부분을 차지하므로, 슬라이딩 윈도우와 같은 기법을 통해 필요한 데이터만 저장하는 것이 중요합니다. 이는 온디바이스 환경에서 모델의 크기와 복잡성을 확장하는 데 필수적입니다.
  • 기능과 성능의 균형: 샘플링과 같은 고급 기능을 도입할 때, 성능 저하를 최소화하는 것이 중요합니다. MLX의 지연 그래프를 활용하여 기존 파이프라인과의 통합을 용이하게 하고, 불필요한 연산을 제거하는 방식으로 이를 달성할 수 있습니다.
  • 점진적 개선: 이 PR은 기존의 torch_native 백엔드에서 발생하던 문제점(예: gpt-oss 모델의 잘못된 서빙)을 해결하고, MLX 백엔드의 기능을 확장하는 방식으로 점진적인 개선을 이루었습니다.

리뷰 댓글 분석

리뷰 댓글은 주로 CI 상태 확인 및 알려진 CI 실패에 대한 병합에 관한 내용이었습니다. 예를 들어 [alexnails] /tag-and-rerun-ci 또는 [alexnails] metal-profiler is known CI failure. merging과 같은 코멘트는 실제 코드 변경 사항에 대한 깊은 기술적 논의보다는 CI/CD 파이프라인 관리와 관련된 것이었습니다. 따라서 이 PR의 코드 변경 사항 자체에 대한 직접적인 기술적 피드백은 리뷰 댓글에서 두드러지지 않았습니다. PR 설명과 코드 diff 자체에 집중하여 분석을 진행했습니다.

References

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글