LLM inference में nondeterminism को कैसे हराया जाए

• LLM (Large Language Model) inference में एक nondeterminism समस्या होती है, जिसमें एक ही input और condition होने पर भी अलग-अलग result आ सकते हैं
• पहले concurrency और floating-point calculation की non-associativity को nondeterminism का मुख्य कारण माना जाता था
• लेकिन वास्तविक निर्णायक कारण batch size बदलने पर kernel (operation code) के अंदर computation order का बदलना है
• अगर kernel के भीतर सभी operation को batch invariance के साथ implement किया जाए, तो पूरी reproducibility की गारंटी दी जा सकती है
• data-parallel operation, split reduction, fixed-size split strategy आदि के जरिए मुख्य operations (RMSNorm, matmul, attention) के लिए batch-invariant kernel बनाए जा सकते हैं

---

परिचय और समस्या का सार

• reproducibility, जो वैज्ञानिक प्रगति का एक मुख्य तत्व है, LLM inference में ठीक से बनी नहीं रहती
• ChatGPT से एक ही सवाल कई बार पूछने पर अक्सर अलग-अलग जवाब बनते हैं
• ऐसा इसलिए होता है क्योंकि LLM में result sampling की प्रक्रिया probability distribution पर आधारित probabilistic selection से होती है
• लेकिन temperature को 0 पर सेट करने के बाद भी व्यवहार में LLM API ज़रूरी नहीं कि deterministic हो, यानी एक ही input पर हर बार वही output आए, यह तय नहीं है
• open source inference libraries (vLLM, SGLang आदि) और अपने hardware पर चलाने पर भी nondeterminism की समस्या बनी रहती है

पुरानी परिकल्पनाएँ और उनकी सीमाएँ

• प्रचलित परिकल्पना: concurrency + floating-point की non-associativity के कारण nondeterminism होता है
  • GPU पर floating-point calculation में operation order और thread completion order के अनुसार result थोड़ा बदल सकता है
• लेकिन वास्तव में एक ही data पर एक ही तरीके से matrix multiplication बार-बार चलाने पर भी हमेशा वही (bitwise equal) result मिलता है
• असली कारण समझने के लिए और गहरा विश्लेषण ज़रूरी है

LLM inference में nondeterminism के कारण का गहन विश्लेषण

floating-point non-associativity का मूल स्वभाव

• floating-point operation में (a+b)+c ≠ a+(b+c) संबंध हो सकता है
• अलग-अलग magnitude (exponent) वाले values पर operation करने पर precision loss, information loss होता है और operation order के अनुसार result बदल सकता है
• क्योंकि operation order बदल सकता है, कई बार sums को random order में करने पर अलग-अलग result निकलते हैं, जिसे प्रयोग में भी देखा गया

kernel के operation order में बदलाव और concurrency

• आम तौर पर atomic add जैसी concurrency समस्याओं को nondeterminism का मुख्य कारण बताया जाता है
• लेकिन LLM inference में इस्तेमाल होने वाले ज़्यादातर kernels, खासकर forward pass में, atomic add के बिना भी काम करते हैं
• पहले से सही parallel strategy, split (reduction) आदि के जरिए एक जैसा operation order सुनिश्चित किया जा सकता है

असली मुख्य कारण: batch invariance की कमी

• हर individual kernel, अगर input वही हो, तो हमेशा वही result देता है (run-to-run deterministic)
• लेकिन कई users के simultaneous request का batch size nondeterministically बदलता रहता है, इसलिए हर request के लिए व्यावहारिक रूप से output स्थिर नहीं रहता
• batch size के अनुसार अंदरूनी रूप से operation को split या merge करने का order बदल जाता है, जिससे nondeterminism पैदा होता है
• यानी server load और parallelism (batch size) का nondeterministic होना ही असली मुख्य कारण है

batch-invariant kernel design और मुख्य operations के उदाहरण

RMSNorm

• data-parallel strategy लागू की जाती है: हर batch element को स्वतंत्र रूप से एक core process करता है
• batch size बड़ा हो तो पर्याप्त parallelism बना रहता है, यानी parallel strategy स्थिर रहती है और batch invariance सुनिश्चित होती है
• batch size बहुत छोटा होने पर split reduction जैसी वैकल्पिक strategy इस्तेमाल की जाती है, लेकिन इस स्थिति में कुछ batch invariance छोड़नी पड़ सकती है

matrix multiplication (matmul)

• tile के आधार पर parallelize करके data-parallel strategy का उपयोग किया जाता है
• Tensor Core optimization के लिए 2D tile split करना पड़ता है, और बहुत छोटे batch में split-K जैसी विशेष strategy की ज़रूरत होती है
• split-K strategy इस्तेमाल करने पर batch invariance टूट सकती है
• performance का कुछ हिस्सा छोड़कर भी एक जैसा kernel configuration मजबूर किया जा सकता है, जिससे स्थिर और reproducible operation order मिलता है

attention

• FlashAttention2 आदि में query-direction parallelization, Key/Value simultaneous reduction strategy से batch invariance सुनिश्चित की जा सकती है
• batch size और sequence split (chunked prefill, prefix caching आदि) के अनुसार reduction order बदल जाए तो invariance टूट जाती है
• split-KV (FlashDecoding) जैसे split-reduction strategy में split size को fixed रखकर operation order को समान रखा जाता है
• अंदरूनी कामकाज में key/value cache और नए token को अलग-अलग process न करके, हर operation में key/value layout को एक जैसा रखना चाहिए

implementation

• vLLM और torch.Library का उपयोग करके batch-invariant kernels लागू किए गए deterministic inference demo उपलब्ध है
• संबंधित replacement kernels GitHub repo (thinking-machines-lab/batch-invariant-ops) में देखे जा सकते हैं

प्रयोग और performance

nondeterminism measurement experiment

• Qwen/Qwen3-235B-A22B-Instruct-2507 model पर temperature 0 की condition में वही prompt (“Tell me about Richard Feynman”) 1000 बार generate किया गया
• 80 अलग-अलग completions बने, यानी same prompt पर भी nondeterminism मौजूद था
• पहले 102 tokens तक result समान थे, और 103वें token पर पहली branching हुई (“Queens, New York” vs “New York City”)
• batch-invariant kernels इस्तेमाल करने पर 1000ों बार बिल्कुल वही result मिला, यानी पूरी reproducibility हासिल हुई

performance evaluation

• 1 GPU पर Qwen-3-8B चलाकर, 90~110 लंबाई वाली sequence के 1000 request चलाए गए
  • default vLLM: 26 सेकंड
  • unoptimized deterministic vLLM: 55 सेकंड
  • improved attention kernel लागू करने पर: 42 सेकंड
• optimization अभी पर्याप्त नहीं है, लेकिन व्यावहारिक रूप से उपयोगी performance स्तर बना रहता है

on-policy RL में महत्व

• पहले training और inference के बीच सूक्ष्म numerical difference के कारण on-policy RL को सटीक रूप से implement नहीं किया जा सकता था
• deterministic inference संभव होने पर sampling और training दोनों को bitwise identical बनाया जा सकता है, जिससे वास्तविक on-policy RL implement करना संभव होता है
• KL-divergence, reward जैसे मुख्य metrics पर पूरी तरह मेल खाते result देखे गए

निष्कर्ष

• LLM inference systems में nondeterminism और numerical error को नज़रअंदाज़ करना आसान है, लेकिन इस समस्या के मूल कारण (batch invariance की कमी) को समझकर सुधार लिया जाए तो पूरी reproducibility और determinism हासिल की जा सकती है
• यह शोध LLM inference की nondeterminism समस्या के समाधान का रास्ता दिखाता है, जिससे developers अपनी systems में पूरी reproducibility हासिल कर सकते हैं

citation जानकारी

• इस शोध को cite करते समय नीचे दी गई जानकारी का उपयोग करें

He, Horace and Thinking Machines Lab, "Defeating Nondeterminism in LLM Inference", 
Thinking Machines Lab: Connectionism, Sep 2025.

या

@article{he2025nondeterminism,
  author = {Horace He and Thinking Machines Lab},
  title = {Defeating Nondeterminism in LLM Inference},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/…},
  doi = {10.64434/tml.20250910}
}