Debug OOM pada backend MoE biasanya bukan sekadar masalah “GPU kurang besar”. Dalam banyak kasus, akar masalahnya ada pada interaksi antara dynamic batching, router yang mengarahkan token ke expert secara tidak merata, buffer sementara yang membengkak, dan konkurensi request yang tidak dibatasi dengan ketat. Gejalanya terlihat di produksi: latensi p95/p99 naik, pod restart karena out-of-memory, request timeout, lalu throughput justru turun meski autoscaling aktif.

Artikel ini membahas studi kasus debugging backend inferensi model Mixture of Experts (MoE) skala besar, terinspirasi konteks sistem seperti LongCat-2.0. Fokusnya praktis: bagaimana mengenali pola OOM yang salah diagnosis, metrik apa yang perlu diperiksa, kenapa batch router bisa memicu ledakan memori, dan guardrail apa yang sebaiknya dipasang agar endpoint inferensi tetap stabil.

Ringkasan masalah: kenapa backend MoE mudah OOM

Pada model dense biasa, ukuran aktivasi relatif lebih mudah diperkirakan dari kombinasi batch size, panjang konteks, dan ukuran hidden state. Pada MoE, ada lapisan tambahan: routing token ke expert. Secara teori, hanya sebagian expert yang aktif untuk tiap token. Namun di backend inferensi, implementasi routing ini membawa beberapa risiko operasional:

  • Batch terlalu besar: aggregator menggabungkan terlalu banyak request demi throughput, tetapi memori untuk aktivasi, indeks routing, gather/scatter buffer, dan output sementara ikut membesar.
  • Distribusi token tidak merata: jika banyak token jatuh ke expert yang sama, kapasitas buffer expert tertentu bisa melonjak tajam.
  • Temporary tensor menumpuk: hasil packing token per expert, mask, indeks, dan buffer komunikasi bisa hidup lebih lama dari yang diperkirakan.
  • Konkurensi request: beberapa batch aktif bersamaan, masing-masing membuat buffer sendiri. OOM sering terjadi bukan pada satu request besar, tetapi pada beberapa request “cukup besar” yang overlap.

Efek akhirnya adalah pola memori yang tampak tidak linear. Kenaikan traffic 20% bisa menghasilkan lonjakan memori jauh di atas 20% jika sistem batching dan routing berada di titik yang salah.

Gejala di produksi yang terlihat lebih dulu

Kasus ini biasanya tidak dimulai dari alert OOM, tetapi dari degradasi kualitas layanan. Gejala yang umum:

  • Latensi p95 dan p99 naik tajam, sementara p50 masih terlihat “normal”.
  • Timeout di sisi gateway atau klien meningkat.
  • Pod inference restart sporadis dengan alasan OOMKilled.
  • GPU memory atau host memory tampak “naik-turun”, tetapi tidak kembali ke baseline dengan cepat.
  • Autoscaling menambah replika, tetapi error rate tidak langsung turun.

Kenapa p50 sering terlihat aman? Karena masalah ini cenderung terjadi pada kombinasi request tertentu: konteks panjang, batch yang kebetulan padat, dan distribusi routing yang berat ke beberapa expert. Jadi, agregat rata-rata sering menutupi kondisi terburuk.

Arsitektur singkat jalur inferensi yang relevan

Untuk memahami sumber OOM, cukup fokus pada alur berikut:

  1. Request masuk ke endpoint inferensi.
  2. Backend melakukan queueing singkat untuk membentuk batch.
  3. Token dari beberapa request digabung.
  4. Router MoE menentukan expert untuk tiap token.
  5. Token dipacking per expert, diproses, lalu hasilnya digabung kembali.
  6. Output dikembalikan ke masing-masing request.

Titik rawan memori ada di langkah 2 sampai 5. Bukan hanya tensor utama model yang memakan memori, tetapi juga struktur perantara untuk routing dan reorganisasi token.

Investigasi: metrik dan log yang benar-benar membantu

1. Korelasikan OOM dengan ukuran batch nyata

Jangan hanya melihat configured max batch size. Catat ukuran batch aktual per eksekusi, misalnya:

  • jumlah request per batch,
  • total token input per batch,
  • perkiraan token decode aktif,
  • panjang konteks request terpanjang,
  • jumlah batch aktif bersamaan.

Dalam banyak sistem, OOM terjadi saat dua atau tiga batch besar overlap, bukan saat satu batch mencapai batas maksimum.

2. Log distribusi routing per expert

Rata-rata top-k routing yang terlihat “seimbang” tidak cukup. Yang perlu dicatat adalah:

  • jumlah token per expert per batch,
  • expert terpadat dalam satu batch,
  • rasio expert terpadat terhadap median expert,
  • jumlah token yang terpaksa dipadding atau di-buffer ulang.

Jika satu atau dua expert menerima jauh lebih banyak token dibanding yang lain, memori buffer lokal bisa melonjak meskipun total token batch belum terlihat ekstrem.

3. Bedakan GPU OOM, CPU OOM, dan allocator pressure

Gejala restart pod tidak selalu berarti GPU kehabisan memori lebih dulu. Ada pola lain:

  • GPU OOM: eksekusi kernel gagal, allocator device gagal memenuhi permintaan.
  • CPU OOM: proses mati karena buffer staging, serialisasi, queue, atau pinned memory membesar.
  • Fragmentasi/allocator pressure: secara total masih ada free memory, tetapi blok kontigu yang dibutuhkan tidak tersedia.

Karena itu, simpan metrik memori host, device, pinned memory, dan jumlah tensor sementara yang dibuat per langkah inferensi bila runtime memungkinkan.

4. Lihat konkurensi internal, bukan hanya QPS

QPS yang sama bisa menghasilkan profil memori berbeda jika scheduler internal berubah. Misalnya, endpoint yang sebelumnya hanya mengizinkan satu batch decode aktif per worker kini mengizinkan beberapa batch overlap demi throughput. Tanpa guardrail, perubahan seperti ini dapat menaikkan puncak memori secara tajam.

Hipotesis yang sempat salah

Dalam studi kasus seperti ini, ada beberapa hipotesis yang terdengar masuk akal tetapi ternyata bukan akar masalah utama:

  • “Ini memory leak.” Tidak selalu. Banyak kasus adalah peak memory explosion, bukan kebocoran jangka panjang. Memori naik sangat tinggi pada batch tertentu lalu proses mati sebelum sempat turun.
  • “Tokenizer atau preprocessing bermasalah.” Bisa berkontribusi, tetapi biasanya bukan penyebab lonjakan tajam saat fase komputasi MoE dimulai.
  • “Autoscaling kurang agresif.” Menambah replika membantu antrean global, tetapi tidak memperbaiki pola batch buruk di dalam satu worker.
  • “Semua karena context length panjang.” Konteks panjang memang mahal, tetapi OOM sering baru terjadi ketika konteks panjang bertemu batch router yang terlalu agresif dan distribusi expert yang skewed.

Hipotesis-hipotesis ini perlu diuji cepat dengan data, bukan intuisi.

Temuan akhir: batch router memicu ledakan memori

Akar masalah yang paling sering muncul pada backend MoE adalah kombinasi empat hal berikut.

1. Dynamic batching hanya dibatasi oleh jumlah request, bukan total token

Ini kesalahan desain yang sangat umum. Dua belas request pendek dan dua belas request sangat panjang dianggap sama oleh scheduler, padahal jejak memorinya jauh berbeda. Pada MoE, perbedaan ini makin tajam karena routing dan packing token juga ikut membesar.

Jika router hanya memakai aturan seperti “maksimal 16 request per batch”, maka satu batch bisa berisi total token yang sangat berlebihan tanpa terlihat melanggar konfigurasi.

2. Routing token ke expert tidak merata

MoE bekerja baik ketika distribusi token antar expert cukup terkendali. Dalam praktik inferensi, distribusi dapat skewed untuk pola input tertentu. Ketika banyak token jatuh ke expert yang sama, backend perlu:

  • menyimpan indeks token untuk expert itu,
  • membuat buffer packed tensor yang lebih besar,
  • melakukan gather/scatter lebih berat,
  • menyimpan temporary output sebelum digabung kembali.

Ledakan memori sering terjadi di sini, bukan pada bobot expert itu sendiri.

3. Temporary tensor bertahan lebih lama karena overlap pipeline

Secara teori, buffer sementara bisa dibebaskan setelah langkah tertentu selesai. Namun pada implementasi nyata, overlap antara batching, prefill, decode, dan postprocessing dapat membuat beberapa temporary tensor hidup bersamaan. Jika ditambah konkurensi request, total puncak memori menjadi jauh di atas estimasi per-request.

4. Tidak ada backpressure per worker

Backend menerima request lebih banyak daripada yang aman untuk diproses paralel di worker yang sama. Akibatnya, batching layer terus membangun antrean “menguntungkan throughput”, tetapi secara lokal justru menciptakan beberapa batch mahal yang saling overlap.

Contoh instrumentasi yang sebaiknya ditambahkan

Tanpa instrumentasi yang tepat, debugging OOM pada backend MoE hanya jadi tebak-tebakan. Berikut contoh struktur log yang lebih berguna daripada log generik “CUDA OOM” atau “process killed”.

event=inference_batch_start
worker_id=inf-3
batch_id=b-184233
request_count=9
total_input_tokens=18432
max_request_tokens=4096
active_decode_sequences=9
concurrent_batches=2

event=moe_routing_stats
batch_id=b-184233
layer=router_17
top_expert_tokens=5210
median_expert_tokens=740
expert_skew_ratio=7.04
routing_overflow_tokens=0

event=memory_snapshot
batch_id=b-184233
host_rss_mb=23140
gpu_allocated_mb=67890
gpu_reserved_mb=72100
pinned_mem_mb=1820
stage=post_pack_pre_expert

Anda tidak harus memakai field yang sama persis, tetapi prinsipnya adalah: catat ukuran kerja aktual, bukan hanya konfigurasi. Dari log seperti ini, korelasi antara batch besar, skew expert, dan lonjakan memori akan lebih cepat terlihat.

Perbaikan yang efektif dan realistis

1. Batasi batch berdasarkan token budget, bukan jumlah request

Perbaikan paling penting adalah mengubah admission control dari “maks request per batch” menjadi maks total token per batch, atau lebih aman lagi, token budget terpisah untuk prefill dan decode.

Kenapa ini bekerja? Karena biaya memori lebih dekat dengan jumlah token aktif daripada jumlah request. Pada MoE, token budget juga lebih representatif terhadap ukuran buffer routing.

class BatchBudget:
    def __init__(self, max_requests, max_prefill_tokens, max_decode_tokens):
        self.max_requests = max_requests
        self.max_prefill_tokens = max_prefill_tokens
        self.max_decode_tokens = max_decode_tokens

    def can_add(self, batch, req):
        if len(batch.requests) >= self.max_requests:
            return False
        if batch.prefill_tokens + req.prefill_tokens > self.max_prefill_tokens:
            return False
        if batch.decode_tokens + req.decode_tokens > self.max_decode_tokens:
            return False
        return True

Nilai ambang harus dikalibrasi berdasarkan profil model dan hardware Anda, tetapi konsepnya tetap sama: token-aware batching lebih aman daripada request-count batching.

2. Pasang batas konkurensi internal per worker

Jangan biarkan satu worker menjalankan terlalu banyak batch aktif secara bersamaan. Batas ini bisa lebih konservatif pada fase prefill, karena jejak memorinya biasanya lebih berat dibanding decode token demi token.

Jika scheduler mendukung prioritas, lebih baik menahan request sebentar di antrean daripada memaksakan overlap batch yang mengakibatkan OOM dan restart proses.

3. Tambahkan guardrail terhadap skew routing

Jika observabilitas menunjukkan beberapa batch menghasilkan distribusi expert yang sangat timpang, tambahkan guardrail operasional seperti:

  • menurunkan token budget saat expert skew ratio melewati ambang tertentu,
  • membatasi ukuran batch ketika request dengan konteks sangat panjang ikut masuk,
  • memecah batch besar menjadi sub-batch jika prediksi biaya routing terlalu tinggi.

Ini memang bisa menurunkan throughput puncak, tetapi jauh lebih murah daripada OOM dan restart berulang.

4. Kurangi lifetime temporary buffer

Tinjau ulang jalur implementasi yang membuat tensor sementara bertahan lebih lama dari perlu. Beberapa perbaikan yang sering membantu:

  • hindari menyimpan hasil intermediate yang sebenarnya bisa langsung dipakai lalu dilepas,
  • pastikan objek referensi lama tidak tertahan oleh closure, cache, atau struktur debug,
  • pisahkan tahap logging agar tidak menahan tensor besar di memori,
  • jangan menyalin tensor ke host kecuali benar-benar dibutuhkan untuk observabilitas.

Di sistem Python, referensi yang tidak sengaja tertahan sering membuat memori terlihat seperti bocor, padahal penyebabnya adalah siklus hidup objek yang terlalu panjang.

5. Terapkan backpressure di endpoint inferensi

Endpoint inferensi harus berani berkata “tunggu” atau “tolak” saat kapasitas internal tidak aman. Praktiknya bisa berupa:

  • batas antrean per worker,
  • timeout antrean yang eksplisit,
  • HTTP 429 atau status overload yang jelas,
  • pengalihan ke replika lain jika scheduler global mendukung.

Tanpa backpressure, sistem cenderung memindahkan masalah dari antrean ke memori, dan itu hampir selalu berakhir lebih buruk.

Checklist mitigasi untuk developer backend

  • Gunakan token-based batching, bukan hanya request-based batching.
  • Pisahkan budget prefill dan decode.
  • Batasi jumlah batch aktif per worker.
  • Log total token aktual, bukan hanya ukuran batch nominal.
  • Log distribusi token per expert dan rasio skew.
  • Monitor host RSS, device allocated/reserved, dan pinned memory.
  • Pasang batas antrean internal dan timeout antrean.
  • Turunkan batch aggressiveness saat konteks panjang mendominasi trafik.
  • Audit lifetime tensor sementara dan referensi debug.
  • Uji skenario burst traffic dengan campuran request pendek dan panjang, bukan hanya beban rata-rata.

Guardrail konfigurasi yang layak dipasang

Guardrail berikut tidak bergantung pada framework tertentu dan bisa diterapkan di banyak backend inferensi:

1. Admission control berbasis biaya

  • Maks total token input per batch.
  • Maks token decode aktif per worker.
  • Maks request panjang per batch.

2. Soft limit dan hard limit memori

  • Soft limit: saat terlewati, scheduler menurunkan ukuran batch.
  • Hard limit: saat mendekat, request baru tidak diterima worker tersebut.

Soft limit penting agar sistem sempat mengerem sebelum kernel atau proses benar-benar gagal.

3. Overload mode

Sediakan mode darurat saat traffic atau pola request buruk muncul:

  • matikan batching agresif sementara,
  • kurangi konkurensi decode,
  • tolak request dengan konteks melebihi ambang sementara,
  • prioritaskan request interaktif daripada batch offline.

Desain endpoint inferensi yang lebih stabil

Pelajaran terbesar dari kasus Debug OOM pada backend MoE adalah bahwa endpoint inferensi tidak boleh didesain seperti API biasa yang hanya memikirkan QPS. Endpoint ini harus sadar biaya komputasi dan memori dari setiap request.

Prinsip desain yang biasanya lebih stabil:

  • Cost-aware scheduling: keputusan batching berdasarkan token, konteks, dan kapasitas worker.
  • Isolation: pisahkan jalur request berat dan ringan bila pola traffic sangat berbeda.
  • Backpressure first: lebih baik menahan sebagian request daripada membiarkan worker kolaps.
  • Observability per tahap: prefill, routing, expert execution, decode, dan postprocess harus punya metrik masing-masing.

Jika semua request diperlakukan sama, backend MoE akan sering terlihat “stabil di rata-rata, rapuh di puncak”.

Langkah uji setelah perbaikan

Setelah guardrail diterapkan, jangan berhenti pada “pod tidak restart lagi”. Verifikasi beberapa hal berikut:

  1. Latensi p95/p99 membaik, bukan hanya error rate turun.
  2. Worker tidak lagi mengalami puncak memori yang tajam pada batch campuran.
  3. Throughput memang mungkin sedikit turun, tetapi throughput efektif meningkat karena tidak ada restart berulang.
  4. Distribusi ukuran batch menjadi lebih konsisten.
  5. Skew expert ekstrem kini memicu degradasi terkontrol, bukan kegagalan total.

Dengan kata lain, targetnya bukan throughput maksimum sesaat, tetapi throughput stabil di bawah beban nyata.

Penutup

OOM pada layanan inferensi MoE skala besar jarang disebabkan satu bug tunggal. Biasanya ini hasil gabungan dari batching yang terlalu agresif, routing token yang tidak merata, temporary tensor yang menumpuk, dan konkurensi yang tidak dikendalikan. Karena itu, solusinya juga harus sistemik: ubah cara batch dibentuk, pasang backpressure, ukur distribusi expert, dan batasi puncak kerja per worker.

Jika Anda sedang menghadapi latensi naik, pod restart, dan timeout pada backend inferensi MoE, mulai dari tiga hal ini: catat total token aktual per batch, ukur skew routing per expert, dan batasi konkurensi internal. Dalam banyak kasus, tiga langkah tersebut sudah cukup untuk mengubah debugging dari spekulasi menjadi diagnosis yang bisa ditindaklanjuti.