ਜੇਕਰ ਤੁਸੀਂ ਕਦੇ ਕਿਸੇ LLM ਨੂੰ ਲੰਬਾ ਜਵਾਬ ਤਿਆਰ ਕਰਦੇ ਹੋਏ ਦੇਖਿਆ ਹੈ ਅਤੇ ਸੋਚਿਆ ਹੈ ਕਿ ਸ਼ੁਰੂਆਤੀ ਪ੍ਰੋਂਪਟ ਦੀ ਚਮਕ ਤੋਂ ਬਾਅਦ ਇਹ ਇੰਨਾ ਹੌਲੀ ਕਿਉਂ ਹੋ ਜਾਂਦਾ ਹੈ, ਤਾਂ ਤੁਸੀਂ ਅਸਲ ਸਮੇਂ ਵਿੱਚ ਇੱਕ ਹਾਰਡਵੇਅਰ ਬੋਟਲਨੇਕ (hardware bottleneck) ਦੇਖ ਰਹੇ ਹੋ। ਜ਼ਿਆਦਾਤਰ ਡਿਵੈਲਪਰ ਆਪਣੇ Python ਕੋਡ, ਫਰੇਮਵਰਕ, ਜਾਂ ਮਾਡਲ ਦੇ ਵਿਸ਼ਾਲ ਆਕਾਰ ਨੂੰ ਦੋਸ਼ੀ ਮੰਨਦੇ ਹਨ। ਉਹ ਫੰਕਸ਼ਨਾਂ ਦੀ ਪ੍ਰੋਫਾਈਲਿੰਗ ਕਰਦੇ ਹਨ, ਆਪਟੀਮਾਈਜ਼ਰ ਬਦਲਦੇ ਹਨ, ਅਤੇ ਪ੍ਰੀਪ੍ਰੋਸੈਸਿੰਗ ਵਿੱਚੋਂ ਮਿਲੀਸੈਕਿੰਡ ਬਚਾਉਣ ਦੀ ਕੋਸ਼ਿਸ਼ ਕਰਦੇ ਹਨ। ਪਰ ਇਹਨਾਂ ਵਿੱਚੋਂ ਕੁਝ ਵੀ ਅਸਲ ਸਮੱਸਿਆ ਨੂੰ ਹੱਲ ਨਹੀਂ ਕਰਦਾ। ਸਪੀਡ ਦੀ ਸੀਮਾ ਤੁਹਾਡੇ ਸਾਫਟਵੇਅਰ ਵਿੱਚ ਨਹੀਂ ਹੈ। ਇਹ ਸਿਲੀਕਾਨ (silicon) ਵਿੱਚ ਹੈ।
ਹਰ ਲਾਰਜ ਲੈਂਗੂਏਜ ਮਾਡਲ ਇਨਫਰੈਂਸ (inference) ਕੰਮ ਤੁਹਾਡੇ ਸਰਵਰ ਵਿੱਚ ਲੱਗੇ GPU ਦੀਆਂ ਦੋ ਭੌਤਿਕ ਵਿਸ਼ੇਸ਼ਤਾਵਾਂ 'ਤੇ ਨਿਰਭਰ ਕਰਦਾ ਹੈ: ਇਹ ਕਿ ਉਹ ਕਿੰਨੀ ਤੇਜ਼ੀ ਨਾਲ ਗਣਨਾ (numbers crunch) ਕਰ ਸਕਦਾ ਹੈ, ਅਤੇ ਇਹ ਕਿ ਉਹ ਉਹਨਾਂ ਨੰਬਰਾਂ ਨੂੰ ਗਣਨਾ ਲਈ ਸਹੀ ਸਥਾਨ 'ਤੇ ਕਿੰਨੀ ਤੇਜ਼ੀ ਨਾਲ ਲਿਆ ਸਕਦਾ ਹੈ।
ਗਣਨਾ ਸਸਤੀ ਹੈ। ਡੇਟਾ ਨੂੰ ਮੂਵ ਕਰਨਾ ਮਹਿੰਗਾ ਹੈ
GPU ਮਾਰਕੀਟਿੰਗ ਨੂੰ ਕੰਪਿਊਟ (compute) ਬਾਰੇ ਗੱਲ ਕਰਨਾ ਬਹੁਤ ਪਸੰਦ ਹੈ। ਪ੍ਰਤੀ ਸੈਕਿੰਡ ਟ੍ਰਿਲੀਅਨਾਂ ਫਲੋਟਿੰਗ-ਪੁਆਇੰਟ ਕੰਪਿਊਟੇਸ਼ਨਾਂ। ਇਹ ਅੰਕ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ। ਪਰ ਕੰਪਿਊਟ ਸਿਰਫ਼ ਅੱਧੀ ਕਹਾਣੀ ਹੈ। ਅੱਧਾ ਹਿੱਸਾ ਮੈਮੋਰੀ ਬੈਂਡਵਿਡਥ (memory bandwidth) ਹੈ, ਉਹ ਦਰ ਜਿਸ ਨਾਲ ਡੇਟਾ ਹਾਈ-ਬੈਂਡਵਿਡਥ ਮੈਮੋਰੀ ਤੋਂ ਕੰਪਿਊਟ ਕੋਰਜ਼ (compute cores) ਵਿੱਚ ਜਾਂਦਾ ਹੈ ਜਿੱਥੇ ਅਸਲ ਅਰਿਥਮੈਟਿਕ (arithmetic) ਹੁੰਦੀ ਹੈ।
ਇੱਕ LLM ਇਹਨਾਂ ਦੋਵਾਂ ਲਿੰਕਾਂ ਵਿੱਚੋਂ ਜੋ ਕਮਜ਼ੋਰ ਹੈ, ਉਸ ਤੋਂ ਤੇਜ਼ ਨਹੀਂ ਚੱਲ ਸਕਦਾ। ਇੱਕ ਵਪਾਰਕ ਰਸੋਈ ਦੀ ਕਲਪਨਾ ਕਰੋ ਜਿਸ ਵਿੱਚ ਵੀਹ ਮਾਸਟਰ ਸ਼ੈੱਫ ਹਨ। ਭੱਠੀਆਂ ਗਰਮ ਹਨ, ਚਾਕੂ ਤਿੱਖੇ ਹਨ, ਅਤੇ ਹਰ ਰਸੋਈਆ ਤਿਆਰ ਹੈ। ਪਰ ਸਬਜ਼ੀਆਂ ਦੀ ਡਿਲੀਵਰੀ ਸਾਈਕਲ ਰਾਹੀਂ ਆਉਂਦੀ ਹੈ, ਇੱਕ ਵਾਰ ਵਿੱਚ ਇੱਕ ਟੋਕਰੀ। ਰਸੋਈ ਦਾ ਕੰਮ ਰੁਕ ਜਾਂਦਾ ਹੈ। ਹੋਰ ਸ਼ੈੱਫ ਜੋੜਨ ਨਾਲ ਇਸ ਦਾ ਹੱਲ ਨਹੀਂ ਹੋਵੇਗਾ। ਤੇਜ਼ ਭੱਠੀਆਂ ਖਰੀਦਣ ਨਾਲ ਵੀ ਇਸ ਦਾ ਹੱਲ ਨਹੀਂ ਹੋਵੇਗਾ। ਰੁਕਾਵਟ ਸੜਕ ਵਿੱਚ ਹੈ।
ਆਧੁਨਿਕ ਡੇਟਾ ਸੈਂਟਰ GPU ਵਿੱਚ, ਅਰਿਥਮੈਟਿਕ ਯੂਨਿਟਸ ਇੰਨੇ ਸ਼ਕਤੀਸ਼ਾਲੀ ਹਨ ਕਿ ਉਹ ਅਕਸਰ ਆਪਣੀਆਂ ਗਣਨਾਵਾਂ ਖਤਮ ਕਰ ਲੈਂਦੇ ਹਨ ਅਤੇ ਫਿਰ ਵਿਕਸ (weights) ਅਤੇ ਐਕਟੀਵੇਸ਼ਨਜ਼ (activations) ਦੇ ਮੈਮੋਰੀ ਰਾਹੀਂ ਆਉਣ ਦੀ ਉਡੀਕ ਵਿੱਚ ਬੇਕਾਰ ਬੈਠੇ ਰਹਿੰਦੇ ਹਨ। ਇਹ ਅਸੰਤੁਲਨ ਤੁਹਾਡੇ ਕੋਡ ਵਿੱਚ ਕੋਈ ਬੱਗ (bug) ਨਹੀਂ ਹੈ। ਇਹ ਚਿਪਸ ਦੇ ਬਣਨ ਦੀ ਭੌਤਿਕ ਅਸਲੀਅਤ ਹੈ। ਮੈਮੋਰੀ ਬੈਂਡਵਿਡਥ ਰੋਅ (raw) ਕੰਪਿਊਟ ਦੇ ਨਾਲ ਤਾਲਮੇਲ ਨਹੀਂ ਰੱਖ ਸਕੀ ਹੈ, ਅਤੇ LLMs ਇਸ ਅਸੰਤੁਲਨ ਲਈ ਖਾਸ ਤੌਰ 'ਤੇ ਨੁਕਸਾਨਦੇਹ ਹਨ ਕਿਉਂਕਿ ਉਹਨਾਂ ਦੇ ਫਾਰਵਰਡ ਪਾਸ (forward passes) ਲਈ ਹਰ ਇੱਕ ਆਉਟਪੁੱਟ ਟੋਕਨ ਲਈ ਹਰ ਇੱਕ ਪੈਰਾਮੀਟਰ ਨੂੰ ਛੂਹਣਾ ਜ਼ਰੂਰੀ ਹੁੰਦਾ ਹੈ।
ਪ੍ਰੋਂਪਟ ਤੇਜ਼ ਕਿਉਂ ਲੱਗਦੇ ਹਨ ਅਤੇ ਜਨਰੇਸ਼ਨ ਹੌਲੀ ਕਿਉਂ ਲੱਗਦੀ ਹੈ
LLM ਇਨਫਰੈਂਸ ਦੋ ਵੱਖ-ਵੱਖ ਪੜਾਵਾਂ ਵਿੱਚ ਵੰਡਿਆ ਹੁੰਦਾ ਹੈ, ਅਤੇ ਉਹ ਹਾਰਡਵੇਅਰ 'ਤੇ ਬਿਲਕੁਲ ਵੱਖਰੇ ਤਰੀਕੇ ਨਾਲ ਦਬਾਅ ਪਾਉਂਦੇ ਹਨ।
Prefill ਉਦੋਂ ਹੁੰਦਾ ਹੈ ਜਦੋਂ ਤੁਹਾਡਾ ਪ੍ਰੋਂਪਟ ਪਹਿਲੀ ਵਾਰ ਮਾਡਲ ਤੱਕ ਪਹੁੰਚਦਾ ਹੈ। ਸਾਰੇ ਟੋਕਨ ਇਕੱਠੇ ਆਉਂਦੇ ਹਨ। GPU ਵੱਡੇ ਮੈਟ੍ਰਿਕਸ-ਮੈਟ੍ਰਿਕਸ ਮਲਟੀਪਲੀਕੇਸ਼ਨ ਦੀ ਵਰਤੋਂ ਕਰਕੇ ਉਹਨਾਂ ਨੂੰ ਸਹਿਗਤ (parallel) ਤਰੀਕੇ ਨਾਲ ਪ੍ਰੋਸੈਸ ਕਰ ਸਕਦਾ ਹੈ। ਹਜ਼ਾਰਾਂ ਅਰਿਥਮੈਟਿਕ ਯੂਨਿਟਸ ਇੱਕੋ ਸਮੇਂ ਚੱਲਦੇ ਹਨ, ਅਤੇ ਵਰਕਲੋਡ ਸੰਘਣਾ ਰਹਿੰਦਾ ਹੈ। ਇਹ ਪੜਾਅ ਕੰਪਿਊਟ-ਬਾਊਂਡ (compute-bound) ਹੁੰਦਾ ਹੈ। ਸ਼ੁਰੂਆਤ ਵਿੱਚ ਜੋ ਅਚਾਨਕ ਤੇਜ਼ੀ ਤੁਸੀਂ ਦੇਖਦੇ ਹੋ? ਉਹ GPU ਦੁਆਰਾ ਉਹੀ ਕੀਤਾ ਜਾ ਰਿਹਾ ਹੈ ਜਿਸ ਲਈ ਉਸਨੂੰ ਬਣਾਇਆ ਗਿਆ ਸੀ।
Decode ਉਹ ਜਗ੍ਹਾ ਹੈ ਜਿੱਥੇ ਚੀਜ਼ਾਂ ਮੁਸ਼ਕਲ ਹੋ ਜਾਂਦੀਆਂ ਹਨ। ਜਦੋਂ ਮਾਡਲ ਅਗਲਾ ਟੋਕਨ ਤਿਆਰ ਕਰਦਾ ਹੈ, ਤਾਂ ਉਹ ਇਹ ਇੱਕ ਵਾਰ ਵਿੱਚ ਇੱਕ ਟੋਕਨ ਕਰਕੇ ਕਰਦਾ ਹੈ। ਇਹ ਪੜਾਅ ਮੈਟ੍ਰਿਕਸ-ਵੈਕਟਰ ਕੰਪਿਊਟੇਸ਼ਨਾਂ 'ਤੇ ਨਿਰਭਰ ਕਰਦਾ ਹੈ, ਜੋ GPU ਦੀ ਸਹਿਗਤ ਸਮਰੱਥਾ ਦੇ ਸਿਰਫ਼ ਇੱਕ ਬਹੁਤ ਛੋਟੇ ਹਿੱਸੇ ਦੀ ਵਰਤੋਂ ਕਰਦੇ ਹਨ। ਇਸ ਤੋਂ ਵੀ ਮਾੜਾ ਇਹ ਹੈ ਕਿ ਹਰ ਨਵਾਂ ਟੋਕਨ GPU ਨੂੰ ਮੈਮੋਰੀ ਤੋਂ ਪੂਰੇ ਮਾਡਲ ਵਿਕਸ (weights) ਨੂੰ ਦੁਬਾਰਾ ਲੋਡ ਕਰਨ ਲਈ ਮਜਬੂਰ ਕਰਦਾ ਹੈ। ਅਰਿਥਮੈਟਿਕ ਯੂਨਿਟਸ ਕੰਮ ਕਰਨਾ ਚਾਹੁੰਦੇ ਹਨ। ਪਰ ਇਸ ਦੀ ਬਜਾਏ, ਉਹ ਉਡੀਕਦੇ ਹਨ। Decode ਮੈਮੋਰੀ-ਬਾਊਂਡ (memory-bound) ਹੁੰਦਾ ਹੈ। GPU ਅਸਲ ਵਿੱਚ ਇੱਕ ਮਹਿੰਗੇ ਟ੍ਰੈਫਿਕ ਕੰਟਰੋਲਰ ਵਜੋਂ ਕੰਮ ਕਰ ਰਿਹਾ ਹੈ, ਜੋ ਮੈਥ ਇੰਜਣਾਂ ਦੇ ਠੰਡੇ ਹੋਣ ਤੱਕ ਮੈਮੋਰੀ ਬੱਸ ਰਾਹੀਂ ਪੈਰਾਮੀਟਰਾਂ ਨੂੰ ਇੱਧਰ-ਉੱਧਰ ਲੈ ਕੇ ਜਾਂਦਾ ਹੈ। ਇਹੀ ਕਾਰਨ ਹੈ ਕਿ ਇੱਕ ਸੌ ਸ਼ਬਦਾਂ ਦਾ ਜਵਾਬ ਦਸ ਸੈਕਿੰਡ ਲੈ ਸਕਦਾ ਹੈ ਭਾਵੇਂ ਸ਼ੁਰੂਆਤੀ ਪ੍ਰੋਂਪਟ ਵਿਸ਼ਲੇਸ਼ਣ ਤੁਰੰਤ ਲੱਗਿਆ ਹੋਵੇ।
KV ਕੈਸ਼ (cache) ਇਸ ਨੂੰ ਹੋਰ ਵੀ ਦਿਲਚਸਪ ਬਣਾਉਂਦਾ ਹੈ। Decode ਦੌਰਾਨ, ਮਾਡਲ ਹਰ ਪਿਛਲੇ ਟੋਕਨ ਲਈ ਕੀ (key) ਅਤੇ ਵੈਲਯੂ (value) ਟੈਂਸਰਾਂ ਨੂੰ ਸਟੋਰ ਕਰਦਾ ਹੈ ਤਾਂ ਜੋ ਉਹ ਅਟੈਂਸ਼ਨ (attention) ਨੂੰ ਸ਼ੁਰੂ ਤੋਂ ਦੁਬਾਰਾ ਕੰਪਿਊਟ ਨਾ ਕਰੇ। ਉਹ ਕੈਸ਼ ਸੀਕੁਇੰਸ ਲੰਬਾਈ ਦੇ ਨਾਲ ਵਧਦਾ ਹੈ। ਇਹ ਮੈਮੋਰੀ ਵਿੱਚ ਵੀ ਰਹਿੰਦਾ ਹੈ। ਇਸ ਲਈ ਹੁਣ GPU ਸਿਰਫ਼ ਵਿਕਸ (weights) ਨੂੰ ਦੁਬਾਰਾ ਲੋਡ ਨਹੀਂ ਕਰ ਰਿਹਾ; ਇਹ ਹਰ ਇੱਕ ਫਾਰਵਰਡ ਪਾਸ 'ਤੇ ਲਗਾਤਾਰ ਵਧ ਰਹੇ ਕੈਸ਼ ਨੂੰ ਪੜ੍ਹ ਅਤੇ ਲਿਖ ਰਿਹਾ ਹੈ। ਕੰਪਿਊਟ ਕੋਰਜ਼ ਬਹੁਤ ਘੱਟ ਮਿਹਨਤ ਕਰ ਰਹੇ ਹਨ ਜਦੋਂ ਕਿ ਮੈਮੋਰੀ ਬੱਸ ਦੋਵਾਂ ਲਈ ਪਸੀਨਾ ਵਹਾ ਰਹੀ ਹੈ।
ਮੈਮੋਰੀ ਵਾਲ (Memory Wall) ਨਾਲ ਲੜਨਾ
ਇੰਜੀਨੀਅਰਾਂ ਨੇ ਇਸ ਗੱਲ ਨੂੰ ਘਟਾਉਣ ਲਈ ਤਕਨੀਕਾਂ ਦਾ ਇੱਕ ਛੋਟਾ ਹਥਿਆਰ ਵਿਕਸਿਤ ਕੀਤਾ ਹੈ ਕਿ ਕਿੰਨਾ ਡੇਟਾ ਮੂਵ ਕਰਨਾ ਪਵੇਗਾ, ਜਾਂ ਘੱਟੋ-ਘੱਟ ਇਸ ਨੂੰ ਮੂਵ ਕਰਨ ਦੀ ਲਾਗਤ ਨੂੰ ਸਾਂਝਾ ਕਰਨ ਲਈ।
Batching ਸਭ ਤੋਂ ਸਿੱਧਾ ਤਰੀਕਾ ਹੈ। ਜੇਕਰ ਇੱਕ ਉਪਭੋਗਤਾ ਦੀ ਬੇਨਤੀ ਮੈਮੋਰੀ ਤੋਂ ਪੂਰੇ ਵਿਕਸ ਲੋਡ ਨੂੰ ਮਜਬੂਰ ਕਰਦੀ ਹੈ, ਤਾਂ ਇੱਕੋ ਸਮੇਂ ਅੱਠ ਜਾਂ ਸੋਲ਼ਾਂ ਬੇਨਤੀਆਂ ਨੂੰ ਪ੍ਰੋਸੈਸ ਕਰਨ ਨਾਲ GPU ਉਸ ਲੋਡ ਨੂੰ ਉਹਨਾਂ ਸਾਰਿਆਂ ਵਿੱਚ ਵੰਡ ਸਕਦਾ ਹੈ। ਵਿਕਸ ਨੂੰ ਇੱਕ ਵਾਰ ਪੜ੍ਹਿਆ ਜਾਂਦਾ ਹੈ ਅਤੇ ਬੈਚ ਵਿੱਚ ਹਰ ਸੀਕੁਇੰਸ ਲਈ ਦੁਬਾਰਾ ਵਰਤਿਆ ਜਾਂਦਾ ਹੈ। ਪ੍ਰੋਡਕਸ਼ਨ ਵਿੱਚ, ਸੋਝੀ ਭਰਪੂਰ ਸ਼ਡਿਊਲਿੰਗ ਪ੍ਰਣਾਲੀਆਂ ਬੇਨਤੀਆਂ ਨੂੰ ਗਤੀਸ਼ੀਲ ਰੂਪ ਵਿੱਚ ਸਮੂਹ ਬਣਾਉਂਦੀਆਂ ਹਨ, ਜਿਸ ਨੂੰ ਕਦੇ-ਕਦੇ continuous ਜਾਂ in-flight batching ਕਿਹਾ ਜਾਂਦਾ ਹੈ, ਤਾਂ ਜੋ GPU ਬਹੁਤ ਘੱਟ ਰੁਕਦਾ ਹੈ। ਇਹ ਇੱਕੋ ਰੂਟ 'ਤੇ ਇੱਕ ਬੱਸ ਅਤੇ ਸੋਲ਼ह ਵੱਖ-ਵੱਖ ਕਾਰਾਂ ਦੇ ਵਿਚਕਾਰਲੇ ਅੰਤਰ ਵਰਗਾ ਹੈ।
Quantization ਬੈਂਡਵਿਡਥ ਦੀ ਸਮੱਸਿਆ 'ਤੇ ਸਿੱਧਾ ਹਮਲਾ ਕਰਦਾ ਹੈ। ਮਾਡਲ ਵੇਟਸ (weights) ਆਮ ਤੌਰ 'ਤੇ ਸਿੱਖ਼ਰਾਂ-ਬਿਟ ਫਲੋਟਿੰਗ-ਪੁਆਇੰਟ ਫਾਰਮੈਟਾਂ ਵਿੱਚ ਸਟੋਰ ਕੀਤੇ ਜਾਂਦੇ ਹਨ। ਉਹਨਾਂ ਨੂੰ ਅੱਠ-ਬਿਟ ਜਾਂ ਚਾਰ-ਬਿਟ ਇੰਟਿਜਰਾਂ ਤੱਕ ਕੰਪਰੈੱਸ ਕਰਕੇ, ਤੁਸੀਂ ਬੱਸ (bus) ਰਾਹੀਂ ਜਾਣ ਵਾਲੇ ਡੇਟਾ ਦੀ ਮਾਤਰਾ ਨੂੰ ਅੱਧਾ ਜਾਂ ਇਸ ਤੋਂ ਵੀ ਵੱਧ ਘਟਾ ਦਿੰਦੇ ਹੋ। ਮਾਡਲ ਨੂੰ ਇਕਸਾਰ ਆਉਟਪੁੱਟ ਦੇਣ ਲਈ ਲੋੜੀਂਦੀ ਸ਼ੁੱਧਤਾ (precision) ਦੀ ਲੋੜ ਹੁੰਦੀ ਹੈ, ਪਰ ਆਧੁਨਿਕ ਪੋਸਟ-ਟ੍ਰੇਨਿੰਗ ਕਵਾਟਾਈਜ਼ੇਸ਼ਨ (post-training quantization) ਤਰੀਕੇ ਗੁਣਵੱਤਾ ਨੂੰ ਖਰਾਬ ਕੀਤੇ ਬਿਨਾਂ ਮਾਡਲ ਦੇ ਮੈਮੋਰੀ ਫੁੱਟਪ੍ਰਿੰਟ ਨੂੰ ਬਹੁਤ ਜ਼ਿਆਦਾ ਘਟਾ ਸਕਦੇ ਹਨ। ਘੱਟ ਡੇਟਾ ਦਾ ਮਤਲਬ ਹੈ ਮੈਮੋਰੀ ਕੰਟਰੋਲਰ 'ਤੇ ਉਡੀਕਣ ਵਿੱਚ ਲੱਗਣ ਵਾਲਾ ਘੱਟ ਸਮਾਂ।
FlashAttention ਅਟੈਂਸ਼ਨ ਮਕੈਨਿਜ਼ਮ (attention mechanism) ਨੂੰ ਇਸ ਤਰ੍ਹਾਂ ਮੁੜ-ਸੰਗਠਿਤ ਕਰਦਾ ਹੈ ਕਿ ਵਿਚਕਾਰਲੇ ਨਤੀਜੇ GPU ਦੀ ਤੇਜ਼ ਆਨ-ਚਿੱਪ ਮੈਮੋਰੀ ਦੇ ਅੰਦਰ ਰਹਿ ਸਕਣ। ਸਟੈਂਡਰਡ ਅਟੈਂਸ਼ਨ ਨੂੰ ਵੱਡੀਆਂ ਅਟੈਂਸ਼ਨ ਮੈਟ੍ਰਿਜ਼ਾਂ ਨੂੰ ਹੌਲੀ ਬਾਹਰੀ ਮੈਮੋਰੀ ਵਿੱਚ ਲਿਖਣਾ ਪੈਂਦਾ ਸੀ ਅਤੇ ਫਿਰ ਉਹਨਾਂ ਨੂੰ ਦੁਬਾਰਾ ਪੜ੍ਹਨਾ ਪੈਂਦਾ ਸੀ। FlashAttention ਗਣਨਾ ਨੂੰ ਛੋਟੇ ਟਾਈਲਾਂ (tiles) ਵਿੱਚ ਵੰਡ ਦਿੰਦਾ ਹੈ ਜੋ SRAM ਵਿੱਚ ਫਿੱਟ ਹੋ ਸਕਦੇ ਹਨ, ਆਨ-ਚਿੱਪ 'ਤੇ ਸੌਫਟਮੈਕਸ (softmax) ਅਤੇ ਸਕੈਲਿੰਗ (scaling) ਕਦਮਾਂ ਨੂੰ ਪੂਰਾ ਕਰਦਾ ਹੈ, ਅਤੇ ਸਿਰਫ਼ ਅੰਤਿਮ ਆਉਟਪੁੱਟ ਨੂੰ ਹੀ ਹਾਈ-ਬੈਂਡਵਿਡਥ ਮੈਮੋਰੀ ਵਿੱਚ ਲਿਖਦਾ ਹੈ। ਇਹ ਮੁੱਖ ਮੈਮੋਰੀ ਦੇ ਘੱਟ ਰੌਂਡ-ਟ੍ਰਿਪਸ (round trips) ਲਈ ਥੋੜ੍ਹੀ ਵਾਧੂ ਕੰਪਿਊਟਿੰਗ ਦਾ ਵਟਾਂਦਰਾ ਕਰਦਾ ਹੈ, ਜੋ ਕਿ ਲਗਭਗ ਹਮੇਸ਼ਾ ਇੱਕ ਫਾਇਦੇਮੰਦ ਫੈਸਲਾ ਹੁੰਦਾ ਹੈ।
PagedAttention ਮੈਮੋਰੀ ਦੀ ਇੱਕ ਵੱਖਰੀ ਕਿਸਮ ਦੀ ਬਰਬਾਦੀ ਨੂੰ ਹੱਲ ਕਰਦਾ ਹੈ। ਡੀਕੋਡ (decode) ਦੌਰਾਨ, KV ਕੈਸ਼ (cache) ਅਣਪਛਾਤੇ ਤਰੀਕੇ ਨਾਲ ਵਧਦਾ ਹੈ। ਰਵਾਇਤੀ ਪ੍ਰਣਾਲੀਆਂ ਹਰੇਕ ਸੀਕੁਇੰਸ (sequence) ਲਈ ਮੈਮੋਰੀ ਦੇ ਨਿਸ਼ਚਿਤ, ਲਗਾਤਾਰ ਟੁਕੜੇ ਅਲਾਟ ਕਰਦੀਆਂ ਹਨ, ਜਿਸ ਨਾਲ ਕੁਝ ਸੀਕੁਇੰਸ ਜਲਦੀ ਖਤਮ ਹੋਣ ਅਤੇ ਕੁਝ ਵਧਣ ਕਾਰਨ ਵੱਡੇ ਖਾਲੀ ਹਿੱਸੇ ਰਹਿ ਜਾਂਦੇ ਹਨ। PagedAttention ਆਪਰੇਟਿੰਗ ਸਿਸਟਮਾਂ ਤੋਂ ਵਰਚੁਅਲ ਮੈਮੋਰੀ (virtual memory) ਦਾ ਸੰਕਲਪ ਉਧਾਰ ਲੈਂਦਾ ਹੈ। ਇਹ KV ਕੈਸ਼ ਐਂਟਰੀਆਂ ਨੂੰ ਨਿਸ਼ਚਿਤ-ਆਕਾਰ ਦੇ ਬਲਾਕਾਂ ਵਿੱਚ ਸਟੋਰ ਕਰਦਾ ਹੈ ਜਿਨ੍ਹਾਂ ਨੂੰ ਗੈਰ-ਲਗਾਤਾਰ ਅਲਾਟ ਕੀਤਾ ਜਾ ਸਕਦਾ ਹੈ ਅਤੇ ਇੱਕ ਇੰਡਾਇਰੈਕਸ਼ਨ ਟੇਬਲ (indirection table) ਰਾਹੀਂ ਮੈਪ ਕੀਤਾ ਜਾ ਸਕਦਾ ਹੈ। ਇਹ ਮੈਮੋਰੀ ਨੂੰ ਰਿਜ਼ਰਵਡ ਪਰ ਅੱਧੇ-ਖਾਲੀ ਬਫਰਾਂ ਦੇ ਅੰਦਰ ਵਿਹਲੇ ਰਹਿਣ ਤੋਂ ਰੋਕਦਾ ਹੈ ਅਤੇ ਵੱਡੇ ਬੈਚ ਸਾਈਜ਼ਾਂ ਦੀ ਇਜਾਜ਼ਤ ਦਿੰਦਾ ਹੈ, ਜੋ ਕਿ ਮੈਮੋਰੀ ਬੱਸ ਨੂੰ ਫਰੈਗਮੈਂਟੇਸ਼ਨ (fragmentation) ਦੇ ਬੋਝ ਦੀ ਬਜਾਏ ਲਾਭਦਾਇਕ ਕੰਮ ਵਿੱਚ ਰੱਖ ਕੇ ਸਮੁੱਚੀ ਥਰੂਪੁੱਟ (throughput) ਵਿੱਚ ਸੁਧਾਰ ਕਰਦਾ ਹੈ।
ਸਵਾਲ ਬਦਲੋ
ਜਦੋਂ ਲੇਟੈਂਸੀ (latency) ਵਧਦੀ ਹੈ, ਤਾਂ ਬਹੁਤ ਸਾਰੀਆਂ ਟੀਮਾਂ ਇਹ ਪੁੱਛਦੀਆਂ ਹਨ ਕਿ ਕੀ ਉਹਨਾਂ ਨੂੰ ਛੋਟੇ ਮਾਡਲ 'ਤੇ ਜਾਣਾ ਚਾਹੀਦਾ ਹੈ ਜਾਂ ਆਪਣੇ ਇਨਫਰੈਂਸ ਸਰਵਰ (inference server) ਨੂੰ ਦੁਬਾਰਾ ਲਿਖਣਾ ਚਾਹੀਦਾ ਹੈ। ਉਹ ਸਵਾਲ ਮਹੱਤਵਪੂਰਨ ਹਨ, ਪਰ ਉਹ ਸੈਕੰਡਰੀ ਹਨ। ਪਹਿਲਾ ਸਵਾਲ ਹਾਰਡਵੇਅਰ ਬਾਰੇ ਹੋਣਾ ਚਾਹੀਦਾ ਹੈ। ਕੀ ਤੁਹਾਡਾ GPU ਅਸਲ ਵਿੱਚ ਕੰਪਿਊਟਿੰਗ ਵਿੱਚ ਰੁੱਝਿਆ ਹੋਇਆ ਹੈ, ਜਾਂ ਇਹ ਡੇਟਾ ਦੀ ਘਾਟ ਨਾਲ ਤੜਫ ਰਿਹਾ ਹੈ?
ਆਪਣੇ ਯੂਟੀਲਾਈਜ਼ੇਸ਼ਨ ਮੈਟ੍ਰਿਕਸ (utilization metrics) ਨੂੰ ਦੇਖੋ। GPU ਕੰਪਿਊਟ ਆਕੂਪੈਂਸੀ (compute occupancy) ਦੇ ਨਾਲ ਮੈਮੋਰੀ ਬੈਂਡਵਿਡਥ ਸੈਚੁਰੇਸ਼ਨ (bandwidth saturation) ਦਾ ਪ੍ਰੋਫਾਈਲ ਕਰੋ। ਜੇਕਰ ਤੁਸੀਂ ਡੀਕੋਡ ਦੌਰਾਨ ਉੱਚ ਮੈਮੋਰੀ ਟਕਰਾਅ (memory contention) ਅਤੇ ਘੱਟ ਅਰਿਥਮੈਟਿਕ ਇੰਟੈਂਸਿਟੀ (arithmetic intensity) ਦੇਖਦੇ ਹੋ, ਤਾਂ ਤੁਹਾਡੇ ਕੋਲ ਮਾਡਲ ਆਰਕੀਟੈਕਚਰ ਦੀ ਸਮੱਸਿਆ ਨਹੀਂ ਹੈ। ਤੁਹਾਡੇ ਕੋਲ ਭੌਤਿਕ ਵਿਗਿਆਨ (physics) ਦੀ ਸਮੱਸਿਆ ਹੈ। ਇਸ ਦਾ ਹੱਲ ਸਾਫ਼-ਸੁਥਰੇ Python ਤੋਂ ਨਹੀਂ ਆਵੇਗਾ। ਇਹ ਵਧੇਰੇ ਹਮਲਾਵਰ ਤਰੀਕੇ ਨਾਲ ਬੈਚਿੰਗ (batching) ਕਰਨ, ਪਾਈਪ ਰਾਹੀਂ ਤੇਜ਼