{ "source_file": "0930", "source_platform": "Medium", "author": "Chaim Rand", "title_listed": "Optimizing Token Generation in PyTorch Decoder Models", "keywords": "", "domain": "大模型算法", "content_type": "深度技术解析", "is_paid_flag": "否", "url": "https://towardsdatascience.com/optimizing-token-generation-in-pytorch-decoder-models/", "publish_date": "2026-02-24", "corpus_stage": "", "notes": "", "id": "081c2758d42a", "collected_at": "2026-10-02T10:56:53", "title": "Optimizing Token Generation in PyTorch Decoder Models | Towards Data Science", "body_markdown": "![Revisit consent button](../images/body-df29beacd7496220cc94.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\n# Optimizing Token Generation in PyTorch Decoder Models\nHiding host-device synchronization via CUDA stream interleaving\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nFebruary 24, 202615 min read\n![](../images/000.jpeg)\n[Photo](https://unsplash.com/photos/time-lapse-photography-of-river--qrcOR33ErA) by Hendrik Cornelissen on Unsplash\nUnderlying the large language models that have pervaded nearly every facet of our daily lives are autoregressive decoder models. These models apply compute-heavy kernel operations to churn out tokens one by one in a manner that, at first glance, seems extremely inefficient. Given the enormous demand for generative AI, it is no surprise that extraordinary engineering effort is being invested into its optimization. Whether it be through custom CUDA kernels, CUDA Graphs, dedicated AI accelerators, or [speculative sampling ](https://medium.com/@xiaxiami/speculative-sampling-in-llms-speeding-up-inference-with-drafts-verification-parallelism-6d948d268a87)— any technique that reduces latency and/or cost by even a fraction of a percentage is a win.\nIn this post, we demonstrate a technique for optimizing token generation in PyTorch using CUDA stream interleaving. While simple to implement, the method addresses a specific, often overlooked bottleneck and can lead to meaningful performance boosts. While pipelining model execution using CUDA streams is common in AI systems engineering, we did not find any tutorial documenting the specific PyTorch-level application we describe here. If you find the technique useful, please be so kind as to reference this post.\nTo facilitate our discussion, we will use a simple GPT-2 PyTorch decoder model from HuggingFace’s transformers (v5.1.0) library. We will run our experiments on an [NVIDIA L40S](https://www.nvidia.com/en-eu/data-center/l40s/) GPU and [PyTorch (2.10.0)](https://pytorch.org/get-started/locally/).\n**Disclaimer:** The code we will share is intended for demonstrative purposes. Please do not rely on its accuracy or optimality. Please do not interpret our mentions of any library, platform, or service as an endorsement of its use.\nImportantly, the value of the CUDA stream-based method we will discuss can vary greatly based on the details of your model and runtime environment. Please be sure to run your own benchmarks before integrating its use.\nOur focus in this post is on PyTorch-native inference workloads which remain extremely prevalent in development and test settings. However, it is important to note that for production environments dedicated LLM inference libraries such as vLLM or NVIDIA TensorRT-LLM tend to deliver greater performance and should be used whenever relevant.\nLearn this step by step with the interactive [AI Engineer](https://roadmap.sh/ai-engineer) roadmap.\n## A Toy GPT-2 Model\nTo simplify our discussion, we will use a GPT-2 decoder model from the HuggingFace `transformers` library and have it run autoregressively on a batch of empty prompts.\nIn the following code block, we initialize the model and define a naive token generation function that creates a batch of random streams up to a given length.\npython\n\n```\nimport torch\nfrom transformers import GPT2LMHeadModel, GPT2Config\n\ntorch.set_float32_matmul_precision('high')\n\nDEVICE = \"cuda\"\n\n# define the decoder model\nconfig = GPT2Config.from_pretrained(\"gpt2\")\nmodel = GPT2LMHeadModel(config).to(DEVICE).eval()\n\n\n@torch.inference_mode()\ndef generate_sequence(model, max_seqlen, batch_size):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n \n for i in range(max_seqlen):\n outputs = model(all_tokens)\n # extract new token\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n # append new token to sequence\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nNext, we define a simple benchmarking function which we use to measure the runtime performance and memory utilization of our token generator in different scenarios.\npython\n\n```\nimport time, statistics\n\n\ndef benchmark(func, num_runs=10):\n # Warmup\n func()\n torch.cuda.synchronize()\n \n runtimes = []\n \n for _ in range(num_runs):\n # reset memory stats before each run\n torch.cuda.empty_cache()\n torch.cuda.reset_peak_memory_stats()\n torch.cuda.synchronize()\n \n start = time.perf_counter()\n _ = func()\n torch.cuda.synchronize()\n end = time.perf_counter()\n \n runtimes.append(end - start)\n \n # Get memory allocator stats from last run\n mem_stats = torch.cuda.memory_stats()\n allocated_peak = mem_stats.get('allocated_bytes.all.peak', 0)\n reserved_peak = mem_stats.get('reserved_bytes.all.peak', 0)\n f_peak = reserved_peak - allocated_peak\n f_pct = (\n 100 * f_peak / reserved_peak\n if reserved_peak > 0 else 0\n )\n \n print(f\"\\n{'='*60}\")\n print(f\"Runtime Results:\")\n print(f\" Mean: {statistics.mean(runtimes):.4f}s\")\n print(f\" Std: {statistics.stdev(runtimes):.4f}s\")\n print(f\" Min: {min(runtimes):.4f}s\")\n print(f\" Max: {max(runtimes):.4f}s\")\n\n print(f\"\\nMemory Stats:\")\n print(f\" Allocated bytes (peak): {allocated_peak / 1e9:.3f} GB\")\n print(f\" Reserved bytes (peak): {reserved_peak / 1e9:.3f} GB\")\n print(f\" Fragmentation (peak): {f_peak / 1e9:.3f} GB ({f_pct:.1f}%)\")\n print(f\"{'='*60}\\n\")\n\n\nbatch_size = 32\nfor max_seqlen in [100, 200, 400]:\n print(\n f\"Benchmarking generation with batch size {batch_size} \"\n f\"and max sequence length {max_seqlen}...\"\n )\n benchmark(\n lambda: generate_sequence(\n model, max_seqlen=max_seqlen, batch_size=batch_size\n )\n )\n```\n\nIn the table below we capture the results for a batch size of 32 and several different sequence lengths:\n![](../images/001.webp)\nBaseline Results (By Author)\nAs the sequence length doubles, the runtime quadruples — appearing to follow a classic _O(N²)_ scaling pattern. Additionally, high memory fragmentation points to severe strain on the CUDA memory allocator, which can result in frequent memory faults and degrade runtime performance. The fragmentation results from each step asking for slightly larger tensor allocations, a pattern which ends up leaving multiple pockets of unusable memory.\nOur first optimization, KV caching, addresses the runtime complexity of our decoder model.\n### KV Caching\nOur naive generator is extremely inefficient — rather than storing and reusing the intermediate tensors from previous tokens, it recalculates the entire sequence at every step.\nWe address the computation inefficiency by using KV caching: We store and reuse the intermediate Key and Value tensors for previous tokens. KV caching reduces the runtime complexity of token generation from _O(N²)_ to _O(N)_.\nIn the following code block, we utilize the transformers library’s built-in support for KV caching to reprogram our token generation function to compute a single batch of tokens in each step.\npython\n\n```\n@torch.inference_mode()\ndef generate_sequence(model, max_seqlen, batch_size, use_cache=False):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n\n # past_key_values is used to store the cached key/values for each layer\n past_key_values = None\n\n for i in range(max_seqlen):\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n use_cache=use_cache\n )\n # update cache for next step\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n # append new token to sequence\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nThe resulting performance numbers are captured in the following table:\n![](../images/body-08e7448ed118a0c17ea2.webp)\nToken Generation With KV Caching (By Author)\nThe performance improvement is profound and, as expected, increases as a function of the sequence length.\nAlthough somewhat better than in our baseline experiment, the degree of memory fragmentation remains a concern. To address this we explore two methods, expandable memory allocations and static KV caching.\n### Expandable CUDA Memory Allocations\nTo reduce CUDA memory fragmentation, we program PyTorch to use [expandable memory segments](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf). As of the time of this writing, this memory optimization is an experimental feature and should be used with caution. Please see [the PyTorch documentation](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf) for details. To use the feature we set the following environment variable:\nbash\n\n```\nexport PYTORCH_ALLOC_CONF=\"expandable_segments:True\"\n```\n\nRerunning our benchmark results in the following table:\n![](../images/body-4e4b721cbe74d2c51ae2.webp)\nKV Caching With Expandable Memory Segments (By Author)\nNot only do we see a marked improvement in fragmentation, but we also get an additional (marginal) improvement in runtime performance.\n### KV Caching With StaticCache\nThe [default cache](https://huggingface.co/docs/transformers/en/kv_cache#default-cache) in HuggingFace is dynamic — it grows as the number of keys and values increases during the generation progresses. HuggingFace supports a fixed-size cache, [StaticCache](https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils#transformers.StaticCache), which pre-allocates a maximum cache size for the KV pairs and reduces strain on the CUDA memory allocator. The disadvantage of using [StaticCache](https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils#transformers.StaticCache) is that the full length of the cache participates in the attention computation at each token generation step, where irrelevant tokens are masked out. This results in a waste of computation that grows with the sequence length. For example, when generating a sequence of 400 tokens, the attention computation for each token will be run on full _400_ X _400_ -sized tensors.\nIn the code block below we enhance our sequence generator to support the use of a [StaticCache](https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils#transformers.StaticCache):\npython\n\n```\nche:\n\nfrom transformers import StaticCache\n\n@torch.inference_mode()\ndef generate_sequence(\n model, max_seqlen, batch_size, use_cache=False, use_static_cache=False\n):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n \n # Initialize static cache if requested\n if use_cache and use_static_cache:\n past_key_values = StaticCache(\n config=config,\n max_batch_size=batch_size,\n max_cache_len=max_seqlen,\n device=DEVICE,\n dtype=model.dtype\n )\n else:\n past_key_values = None\n \n # Initialize cache position tracking for static cache\n cache_positions = torch.arange(max_seqlen, device=DEVICE)\n \n for i in range(max_seqlen):\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n cache_position = (\n cache_positions[i:i+1] if use_static_cache else None\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n cache_position=cache_position,\n use_cache=use_cache\n )\n # update cache for next step\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n # append new token to sequence\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nThe updated results are captured below:\n![](../images/body-26e1fd4c3df70006476e.webp)\nToken Generation With Static KV Cache (By Author)\nUsing a fixed-sized cache greatly improves memory utilization as indicated by the decrease in memory fragmentation. However, its impact on runtime performance is mixed — for 100 tokens it reduces performance compared to a dynamic cache, whereas for 200 and 400 tokens it boosts performance by 9% and 10%, respectively.\nThere are more advanced methods of implementing attention that optimize for memory utilization without the cost of wasted computation. In a previous post, [Optimizing Transformer Models for Variable-Length Input Sequences](https://towardsdatascience.com/optimizing-transformer-models-for-variable-length-input-sequences-19fb88fddf71/), we covered some PyTorch techniques for computing attention _sparsely_ to reduce computation waste. For production settings, libraries such as vLLM use [PagedAttention](https://blog.vllm.ai/2023/06/20/vllm.html) for maximizing memory utilization. These methods are outside the scope of this post.\nFor more details on caching in HuggingFace, please see the [caching strategies overview](https://huggingface.co/docs/transformers/en/kv_cache#cache-strategies).\n### Model Compilation\nOne of the [documented](https://huggingface.co/docs/transformers/en/kv_cache#fixed-size-cache) advantages of using a fixed-sized cache is that it allows for taking advantage of many just-in-time (JIT) optimizations.\nIn the following code block we apply our benchmark to a [PyTorch-compiled](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) version of our decoder model:\npython\n\n```\nbatch_size = 32\nmax_seqlen = 100\n\nmodel = torch.compile(model)\n\nbenchmark(\n lambda: generate_sequence(\n model,\n max_seqlen=max_seqlen,\n batch_size=batch_size,\n use_cache=True,\n use_static_cache=True\n )\n)\n```\n\nModel compilation results in an additional boost to runtime performance as shown in the table below:\n![](../images/body-5473a44dd00f72e3df7d.webp)\nToken Generation With torch.compile (By Author)\nNote that we can apply model compilation when using dynamic caching, as well. However, torch.compile provides the best results when the computation graph is composed of fixed-sized tensors (e.g., see [here](https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/) for more details).\n### The Performance Penalty of Early Stopping\nAn integral part of common token generators is checking for the end-of-sequence (EOS) at the end of each step. Without this test, token generators would always run for _max_seqlen,_ even if all the sequences in the batch have ended. This could result in considerable computation waste and unnecessary latency — especially when common sequence lengths are much shorter than the maximum length. In the case of our toy experiment, we wait for all the sequences in the batch to end and discontinue token generation. Production-grade implementations will commonly perform _continuous batching_ — replacing completed sequences with new prompts on the input queue.\npython\n\n```\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n```\n\nImportantly, the `.item()` call on the `stop_gpu` tensor, triggers a blocking host-device synchronization event. More specifically, in order to evaluate the conditional _if_ statement, the CPU must wait for the GPU to complete its computation and copy the contents of the tensor to host memory. While the CPU waits, it is blocked from executing the next step of the token generation loop, or more accurately, it is blocked from loading the next computation kernels onto the GPU.\nTo measure the impact of the stopping condition on runtime performance, we add instrumentation for performance profiling with [NVIDIA Nsight™ Systems (nsys)](https://docs.nvidia.com/nsight-systems/index.html) using the [torch.cuda.profiler](https://github.com/pytorch/pytorch/blob/main/torch/cuda/profiler.py) and nvtx (v0.2.14) APIs. (See our [recent post](https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/) for more details on performance profiling with nsys).\npython\n\n```\nore details on performance profiling with nsys).\n\nimport nvtx\nfrom torch.cuda import profiler\n\n@torch.inference_mode()\ndef generate_sequence(\n model, max_seqlen, batch_size, use_cache=False, use_static_cache=False\n):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n \n # Initialize static cache if requested\n if use_cache and use_static_cache:\n past_key_values = StaticCache(\n config=config,\n max_batch_size=batch_size,\n max_cache_len=max_seqlen,\n device=DEVICE,\n dtype=model.dtype\n )\n else:\n past_key_values = None\n \n # Initialize cache position tracking for static cache\n cache_positions = torch.arange(max_seqlen, device=DEVICE)\n \n for i in range(max_seqlen):\n if i == 30:\n # start nsys profiler\n torch.cuda.synchronize()\n profiler.start()\n elif i == 50:\n # stop nsys profiler\n torch.cuda.synchronize()\n profiler.stop()\n with nvtx.annotate(f\"Step {i+1}\", color=\"blue\"):\n with nvtx.annotate(\"Model Forward\", color=\"green\"):\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n cache_position = (\n cache_positions[i:i+1] if use_static_cache else None\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n cache_position=cache_position,\n use_cache=use_cache\n )\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n with nvtx.annotate(\"Check Stop Condition\", color=\"red\"):\n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nWe run our script using the _cudaProfilerApi_ option to start and stop the profiler programmatically. Please see the official documentation for full details on [profiling from the nsys CLI](https://docs.nvidia.com/nsight-systems/UserGuide/index.html#profiling-from-the-cli).\nbash\n\n```\nnsys profile \\\n --capture-range=cudaProfilerApi \\\n --trace=cuda,nvtx,osrt \\\n --output=baseline \\\n python train.py\n```\n\nThe following trace, captured for a batch size of 16 and sequence length of 100, shows the GPU idling for about 110 microseconds in between steps — an eternity in the context of high-performance GPU workloads. This is a direct result of the synchronization event triggered by the EOS test.\n![](../images/body-b18938655ed310285221.webp)\nGPU Utilization Drops Between Each Step (By Author)\nIn production-grade implementations such synchronization issues are avoided by some combination of 1) use of lower level (e.g., C/C++) code that avoids the limitation of the Python interpreter, 2) using CUDA graphs to reduce overhead of kernel loading, 3) moving conditional checks onto the GPU using [conditional nodes](https://developer.nvidia.com/blog/dynamic-control-flow-in-cuda-graphs-with-conditional-nodes/), and 4) continuously and asynchronously preparing subsequent requests while the EOS check is in progress.\nIn the next section, we demonstrate a technique for hiding the overhead of the host-device synchronization in PyTorch using [CUDA streams](https://docs.pytorch.org/docs/stable/generated/torch.cuda.Stream_class.html).\n## A CUDA Stream Optimization\nA CUDA stream is a linear sequence of operations (kernels, memory copies, etc.) that execute in order on the GPU. While operations _within_ a single stream are guaranteed to execute sequentially, operations in different streams can execute concurrently or overlap.\nIn previous posts (e.g., [here](https://towardsdatascience.com/pipelining-ai-ml-training-workloads-with-cuda-streams/) and [here](https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/)) we demonstrated the use of CUDA streams in pipelining common AI/ML workloads, e.g., executing a model on batch _N_ while preparing batch _N+1_. In this post we will use CUDA streams to enable the CPU to load the GPU kernels of step _N+1_ before checking the stopping criteria of step _N_. Contrary to our previous demonstrations of CUDA streams, our current example will not necessarily involve concurrent GPU kernel execution. \nWe implement an alternative token generation function that interleaves two CUDA streams, running the following operations iteratively:\nProgram stream _i%2_ to: **(A)** wait for stream (_i-1)%2_ to complete its generation of token _i-1,__**(**_**B)** use the updated tensors to calculate the token _i_ , **(C)** run the EOS test for token _i_ on the GPU, and **(D)** perform a (non-blocking) copy of the EOS test result to pinned memory on the CPU.\nOn the default CUDA stream, wait for stream (_i-1)%2_ to complete its generation of token _i-1._\nOn the default CUDA stream, check if the stopping criteria for token _i-1_ were met. If so, halt the generator and return. Otherwise, increment _i_ and return to step 1.\nWhereas previously, the initialization of token _i_ generation was blocked by the EOS test on token _i-1_ , the use of CUDA streams allows us to program the generation of token _i_ **before** wecheck the result of the EOS test on token _i-1_. In practice, the EOS test for token _i-1_ on the CPU runs while the GPU is computing token _i_.\npython\n\n```\n@torch.inference_mode()\ndef generate_sequence_pipelined(\n model,\n max_seqlen,\n batch_size,\n use_cache=False,\n use_static_cache=False\n):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n past_key_values = None\n \n # Initialize static cache if requested\n if use_cache and use_static_cache:\n past_key_values = StaticCache(\n config=config,\n max_batch_size=batch_size,\n max_cache_len=max_seqlen,\n device=DEVICE,\n dtype=model.dtype\n )\n \n # Initialize cache position tracking for static cache\n cache_positions = torch.arange(max_seqlen, device=DEVICE)\n \n # Dual streams for pipelining\n streams = [torch.cuda.Stream(), torch.cuda.Stream()]\n stop_host = [\n torch.tensor(False, pin_memory=True),\n torch.tensor(False, pin_memory=True)\n ]\n \n for i in range(max_seqlen):\n curr_idx, prev_idx = i % 2, (i+1) % 2\n curr_s, prev_s = streams[curr_idx], streams[prev_idx]\n \n # Launch iteration i in current stream\n with torch.cuda.stream(curr_s):\n # program stream to wait for previous stream to complete\n curr_s.wait_stream(prev_s)\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n cache_position = (\n cache_positions[i:i+1] if use_static_cache else None\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n cache_position=cache_position,\n use_cache=use_cache\n )\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n \n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n stop_host[curr_idx].copy_(stop_gpu, non_blocking=True)\n \n # Check previous iteration's stop signal\n torch.cuda.current_stream().wait_stream(prev_s)\n if stop_host[prev_idx].item():\n print(f\"All sequences finished at step {i}\")\n break\n \n return all_tokens\n```\n\nThe image below captures the nsys trace for our new token generator:\n![](../images/body-fed6bef0164ba922642a.webp)\nConstant GPU Activity When Applying CUDA Streams (By Author)\nIn the CUDA section of the trace we can see the use of two CUDA streams, with token generation being passed back and forth in a sort of ping-pong effect: One stream generates all of the odd tokens and second all of the even tokens. The CPU is about half a step ahead of the GPU — allowing it to program step _i_ while the GPU is computing step _i-1_. The CPU-side EOS stop-check of step _i-1_(in red) occurs after step _i_ is fully programmed (and has started running).Most importantly, we now find the GPU utilization to be consistent — the idling we saw before is gone.\nThe CUDA stream interleaving results in an additional performance boost, as shown in the table below:\n![](../images/body-9dd2d338c040c1f98e47.webp)\nToken Generation With CUDA Streams (By Author)\nWe would expect the benefit of the ping-pong solution we have implemented to be impacted by the ratio between the GPU idle time (i.e., the overhead of kernel loading) and the kernel computation time. To test this, we fix the sequence length at 100 and rerun the benchmark for a number of batch sizes:\n![](../images/body-ad870a33113327e829b9.webp)\nImpact of Pipelining for Varying Batch Size (By Author)\nAs expected, the highest performance gain, 11.6%, occurs when the batch size is smallest and the kernel computation load is at its lowest. As the kernel compute increases, the ratio of kernel loading to kernel compute time decreases as does the impact of CUDA stream interleaving.\nNote that there is some overhead to the use of CUDA streams. This can be demonstrated by comparing our interleaving solution to a token generator that skips the EOS test altogether:\n![](../images/body-ddaf101a5dba6cd682d0.webp)\nOverhead of CUDA Stream Interleaving (By Author)\n### The Potential Performance Pitfalls of Using CUDA Streams\nCUDA streams should be used with extreme caution. When using the [default stream](https://docs.pytorch.org/docs/stable/generated/torch.cuda.default_stream.html) we can rely on PyTorch to perform any necessary synchronization when data is moved around. However, when using CUDA streams, we must ensure appropriate synchronization explicitly. In particular, we must ensure appropriate data transfer between the streams. Otherwise, we may experience CUDA errors (e.g., “device-side assert triggered”) — if we are lucky. If we are less lucky, we may experience data corruption without even knowing it. See the PyTorch [CUDA stream](https://docs.pytorch.org/docs/stable/notes/cuda.html#cuda-streams) documentation for more details on appropriate use.\nFor AI/ML workloads with large CUDA memory utilization, such as LLMs, another consideration is memory utilization. The PyTorch caching allocator manages memory on a per-stream basis; using multiple streams can lead to increased memory reservation and fragmentation. These could result in increased memory faults that might overshadow the potential gains from the use of streams.\n## Results\nIn the table below we summarize the runtime results of applying static caching, compilation, and pipelining on a batch of 32 sequences and a maximum sequence length of 100. The results are sorted in increasing order of performance:\n![](../images/body-656754222c236ea5bc2b.webp)\nToken Generation Optimization Results (By Author)\nIn the case of our toy GPT-2 model, the best results — nearly 5 times the baseline performance — are achieved when employing PyTorch compilation and the CUDA stream interleaving method discussed in this post. However, as we have seen, the impact of CUDA interleaving could vary greatly based on the properties of the workload and runtime environment, particularly on the ratio between the kernel loading time and the kernel compute time. Please be sure to run your own benchmarks before adopting this method.\n## Summary\nIn high-performance AI engineering, any hint of GPU under-utilization presents an opportunity for optimization. One of the primary optimization tools on NVIDIA GPUs is CUDA streams. In this post, we demonstrated their use in solving the idle GPU time that results from the host-device synchronization associated with early-stopping in PyTorch-native autoregressive token generation. By interleaving CUDA streams in a “ping-pong” pattern, we successfully hid the latency imposed by the EOS-check which resulted in a meaningful increase the workload’s throughput. By combining this technique with the well-known methods of model compilation and static caching, we can maximize the performance of PyTorch-native inference.\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](../images/body-97f5483d57a9b8741146.jpeg)\nWritten By\n## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Cuda](https://towardsdatascience.com/tag/cuda/)[Deep Learning](https://towardsdatascience.com/tag/deep-learning/)[Optimization](https://towardsdatascience.com/tag/optimization/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F&title=Optimizing%20Token%20Generation%20in%20PyTorch%20Decoder%20Models)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F&title=Optimizing%20Token%20Generation%20in%20PyTorch%20Decoder%20Models)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F&text=Optimizing%20Token%20Generation%20in%20PyTorch%20Decoder%20Models)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](../images/002.jpeg)](https://towardsdatascience.com/my-story-of-taking-the-tensorflow-developer-certification-exam-85d4fb777cbb/)\n## [My Story of Taking the TensorFlow Developer Certification Exam](https://towardsdatascience.com/my-story-of-taking-the-tensorflow-developer-certification-exam-85d4fb777cbb/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nMy overall experience of taking the exam, how I prepared for it, and what I would've done differently if I had to take the exam again.\n[Ruben Winastwan](https://towardsdatascience.com/author/marcellusruben/)\nJuly 17, 20209 min read\n * [![](../images/003.gif)](https://towardsdatascience.com/towards-infinite-llm-context-windows-e099225abaaf/)\n## [Towards infinite LLM context windows](https://towardsdatascience.com/towards-infinite-llm-context-windows-e099225abaaf/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nFrom 512 to 1M+ tokens in 5 years - LLMs have rapidly expanded their context windows. Where's the limit?\n[Krzysztof K. Zdeb](https://towardsdatascience.com/author/krzysztof.kornel/)\nApril 28, 202410 min read\n * [![](../images/004.png)](https://towardsdatascience.com/resnets-residual-blocks-deep-residual-learning-a231a0ee73d2/)\n## [ResNets - Residual Blocks & Deep Residual Learning](https://towardsdatascience.com/resnets-residual-blocks-deep-residual-learning-a231a0ee73d2/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nDeep Learning harnesses the power of Big Data by building deep neural architectures that try to approximate a function f(x) that can map...\n[Prem Oommen](https://towardsdatascience.com/author/thedataoracle/)\nNovember 28, 20204 min read\n * [![](../images/005.png)](https://towardsdatascience.com/10-examples-of-awful-data-that-i-had-to-work-with-as-a-data-scientist-10894cd35828/)\n## [10 Examples of Awful Data that I had to work with as a Data Scientist](https://towardsdatascience.com/10-examples-of-awful-data-that-i-had-to-work-with-as-a-data-scientist-10894cd35828/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nA glimpse into the frustrating dark side of working with data\n[Terence Shin, MSc, MBA](https://towardsdatascience.com/author/terenceshin/)\nNovember 11, 20206 min read\n * [![](../images/006.jpeg)](https://towardsdatascience.com/outlier-aware-clustering-beyond-k-means-76f7bf8b4899/)\n## [Outlier-Aware Clustering: Beyond K-Means](https://towardsdatascience.com/outlier-aware-clustering-beyond-k-means-76f7bf8b4899/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nCan data science help you choose your next job? First of all, that's a terrible idea.\n[Mathieu Lemay](https://towardsdatascience.com/author/lsci/)\nNovember 29, 20182 min read\n * [![](../images/007.jpg)](https://towardsdatascience.com/will-the-covid-19-vaccine-alter-my-dna-and-hook-me-to-an-ai-f7d086b94e07/)\n## [Will the COVID-19 vaccine alter my DNA and hook me to an AI?](https://towardsdatascience.com/will-the-covid-19-vaccine-alter-my-dna-and-hook-me-to-an-ai-f7d086b94e07/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nChecking your facts in the time of the pandemic\n[Alexiei Dingli](https://towardsdatascience.com/author/alexieidingli/)\nAugust 15, 20205 min read\n * [![](../images/008.jpeg)](https://towardsdatascience.com/how-to-calculate-the-mean-and-standard-deviation-normalizing-datasets-in-pytorch-704bd7d05f4c/)\n## [How to calculate the mean and standard deviation - Normalizing datasets in Pytorch](https://towardsdatascience.com/how-to-calculate-the-mean-and-standard-deviation-normalizing-datasets-in-pytorch-704bd7d05f4c/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nNeural networks converge much faster if the input data is normalized. Learn the reason why and how to implement this.\n[Jorrit Willaert](https://towardsdatascience.com/author/jorritwillaert98/)\nSeptember 24, 20213 min read\n * [![](../images/009.jpeg)](https://towardsdatascience.com/10-examples-to-master-python-dictionary-comprehensions-7aaa536f5960/)\n## [10 Examples to Master Python Dictionary Comprehensions](https://towardsdatascience.com/10-examples-to-master-python-dictionary-comprehensions-7aaa536f5960/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nA comprehensive practical guide\n[Soner Yıldırım](https://towardsdatascience.com/author/sonery/)\nJanuary 2, 20215 min read\n * [![](../images/010.png)](https://towardsdatascience.com/an-existential-crisis-of-a-veteran-researcher-in-the-age-of-generative-ai/)\n## [An Existential Crisis of a Veteran Researcher in the Age of Generative AI](https://towardsdatascience.com/an-existential-crisis-of-a-veteran-researcher-in-the-age-of-generative-ai/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nIf you have had an existential moment in your career due to AI, let me know.\n[Pedram Ataee](https://towardsdatascience.com/author/pedram-ataee/)\nApril 23, 20255 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "metadata": { "title": "Optimizing Token Generation in PyTorch Decoder Models | Towards Data Science", "description": "Hiding host-device synchronization via CUDA stream interleaving", "keywords": null, "author": "Chaim Rand", "og:locale": "en_US", "og:type": "article", "og:title": "Optimizing Token Generation in PyTorch Decoder Models", "og:description": "Hiding host-device synchronization via CUDA stream interleaving", "og:url": "https://towardsdatascience.com/optimizing-token-generation-in-pytorch-decoder-models/", "og:site_name": "Towards Data Science", "og:image": "https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/1-1-1.jpeg", "og:image:width": "1100", "og:image:height": "666", "og:image:type": "image/jpeg", "twitter:card": "summary_large_image", "twitter:site": "@TDataScience", "twitter:creator": "@TDataScience", "twitter:label1": "Written by", "twitter:data1": "Chaim Rand", "twitter:label2": "Est. reading time", "twitter:data2": "15 min read", "article:published_time": "2026-02-24T15:00:07.000Z", "article:modified_time": "2026-02-24T15:00:07.000Z", "article:author": "https://towardsdatascience.com/author/chaimrand/" }, "paywall": { "listed_paid": false, "detected": false, "marker": "", "public_only": false, "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙" }, "crawl": { "success": true, "status_code": 200, "error": "" }, "code_blocks": [ { "language": "txt", "file": "attachments/081c2758d42a/code/block_000.txt", "chars": 1193 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_001.txt", "chars": 1760 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_002.txt", "chars": 1343 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_003.txt", "chars": 53 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_004.txt", "chars": 1925 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_005.txt", "chars": 245 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_006.txt", "chars": 243 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_007.txt", "chars": 2545 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_008.txt", "chars": 118 }, { "language": "txt", "file": "attachments/081c2758d42a/code/block_009.txt", "chars": 2588 } ], "images": [ { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:45:45Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:45:45Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:45:45Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:45:45Z" } ], "audio": [], "videos": [ { "url": "https://www.youtube.com/c/TowardsDataScience", "kind": "channel_reference", "status": "reference_only" } ], "links": [ { "type": "internal", "href": "https://towardsdatascience.com/" }, { "type": "internal", "href": "https://towardsdatascience.com/latest/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/editors-pick/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/the-variable/" }, { "type": "internal", "href": "https://towardsdatascience.com/shipai/" }, { "type": "internal", "href": "https://towardsdatascience.com/submissions/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/chaimrand/" }, { "type": "internal", "href": "https://towardsdatascience.com/optimizing-transformer-models-for-variable-length-input-sequences-19fb88fddf71/" }, { "type": "internal", "href": "https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/" }, { "type": "internal", "href": "https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/" }, { "type": "internal", "href": "https://towardsdatascience.com/pipelining-ai-ml-training-workloads-with-cuda-streams/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/deep-dives/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/cuda/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/deep-learning/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/optimization/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/pytorch/" }, { "type": "internal", "href": "https://towardsdatascience.com/my-story-of-taking-the-tensorflow-developer-certification-exam-85d4fb777cbb/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/deep-learning/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/marcellusruben/" }, { "type": "external", "href": "https://write.insightmediagroup.io" }, { "type": "external", "href": "https://www.linkedin.com/company/towards-data-science/" }, { "type": "external", "href": "https://x.com/TDataScience" }, { "type": "external", "href": "https://unsplash.com/photos/time-lapse-photography-of-river--qrcOR33ErA" }, { "type": "external", "href": "https://medium.com/@xiaxiami/speculative-sampling-in-llms-speeding-up-inference-with-drafts-verification-parallelism-6d948d268a87" }, { "type": "external", "href": "https://www.nvidia.com/en-eu/data-center/l40s/" }, { "type": "external", "href": "https://pytorch.org/get-started/locally/" }, { "type": "external", "href": "https://roadmap.sh/ai-engineer" }, { "type": "external", "href": "https://docs.pytorch.org/docs/stable/notes/cuda.html" }, { "type": "external", "href": "https://huggingface.co/docs/transformers/en/kv_cache" }, { "type": "external", "href": "https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils" }, { "type": "external", "href": "https://blog.vllm.ai/2023/06/20/vllm.html" }, { "type": "external", "href": "https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html" }, { "type": "external", "href": "https://docs.nvidia.com/nsight-systems/index.html" }, { "type": "external", "href": "https://github.com/pytorch/pytorch/blob/main/torch/cuda/profiler.py" }, { "type": "external", "href": "https://docs.nvidia.com/nsight-systems/UserGuide/index.html" }, { "type": "external", "href": "https://developer.nvidia.com/blog/dynamic-control-flow-in-cuda-graphs-with-conditional-nodes/" }, { "type": "external", "href": "https://docs.pytorch.org/docs/stable/generated/torch.cuda.Stream_class.html" }, { "type": "external", "href": "https://docs.pytorch.org/docs/stable/generated/torch.cuda.default_stream.html" }, { "type": "external", "href": "https://www.facebook.com/sharer/sharer.php?title=Optimizing+Token+Generation+in+PyTorch+Decoder+Models&u=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F" } ], "repair_version": 2, "repaired_at": "2026-10-20T17:26:12Z", "source_url": "https://towardsdatascience.com/optimizing-token-generation-in-pytorch-decoder-models/", "delivery_status": "unverified", "discovery_status": "http_403", "missing_reason": "源页面未重新核验:http_403；已按保留正文补采图片", "repair_error": "", "actual_video_files": [], "expected_saved_body_images": 4, "image_discovery": "saved_markdown", "source_body_markdown": "![Revisit consent button](https://cdn-cookieyes.com/assets/images/revisit.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](https://cdn-cookieyes.com/assets/images/close.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\n# Optimizing Token Generation in PyTorch Decoder Models\nHiding host-device synchronization via CUDA stream interleaving\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nFebruary 24, 202615 min read\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/1-1-1.jpeg)\n[Photo](https://unsplash.com/photos/time-lapse-photography-of-river--qrcOR33ErA) by Hendrik Cornelissen on Unsplash\nUnderlying the large language models that have pervaded nearly every facet of our daily lives are autoregressive decoder models. These models apply compute-heavy kernel operations to churn out tokens one by one in a manner that, at first glance, seems extremely inefficient. Given the enormous demand for generative AI, it is no surprise that extraordinary engineering effort is being invested into its optimization. Whether it be through custom CUDA kernels, CUDA Graphs, dedicated AI accelerators, or [speculative sampling ](https://medium.com/@xiaxiami/speculative-sampling-in-llms-speeding-up-inference-with-drafts-verification-parallelism-6d948d268a87)— any technique that reduces latency and/or cost by even a fraction of a percentage is a win.\nIn this post, we demonstrate a technique for optimizing token generation in PyTorch using CUDA stream interleaving. While simple to implement, the method addresses a specific, often overlooked bottleneck and can lead to meaningful performance boosts. While pipelining model execution using CUDA streams is common in AI systems engineering, we did not find any tutorial documenting the specific PyTorch-level application we describe here. If you find the technique useful, please be so kind as to reference this post.\nTo facilitate our discussion, we will use a simple GPT-2 PyTorch decoder model from HuggingFace’s transformers (v5.1.0) library. We will run our experiments on an [NVIDIA L40S](https://www.nvidia.com/en-eu/data-center/l40s/) GPU and [PyTorch (2.10.0)](https://pytorch.org/get-started/locally/).\n**Disclaimer:** The code we will share is intended for demonstrative purposes. Please do not rely on its accuracy or optimality. Please do not interpret our mentions of any library, platform, or service as an endorsement of its use.\nImportantly, the value of the CUDA stream-based method we will discuss can vary greatly based on the details of your model and runtime environment. Please be sure to run your own benchmarks before integrating its use.\nOur focus in this post is on PyTorch-native inference workloads which remain extremely prevalent in development and test settings. However, it is important to note that for production environments dedicated LLM inference libraries such as vLLM or NVIDIA TensorRT-LLM tend to deliver greater performance and should be used whenever relevant.\nLearn this step by step with the interactive [AI Engineer](https://roadmap.sh/ai-engineer) roadmap.\n## A Toy GPT-2 Model\nTo simplify our discussion, we will use a GPT-2 decoder model from the HuggingFace `transformers` library and have it run autoregressively on a batch of empty prompts.\nIn the following code block, we initialize the model and define a naive token generation function that creates a batch of random streams up to a given length.\npython\n\n```\nimport torch\nfrom transformers import GPT2LMHeadModel, GPT2Config\n\ntorch.set_float32_matmul_precision('high')\n\nDEVICE = \"cuda\"\n\n# define the decoder model\nconfig = GPT2Config.from_pretrained(\"gpt2\")\nmodel = GPT2LMHeadModel(config).to(DEVICE).eval()\n\n\n@torch.inference_mode()\ndef generate_sequence(model, max_seqlen, batch_size):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n \n for i in range(max_seqlen):\n outputs = model(all_tokens)\n # extract new token\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n # append new token to sequence\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nNext, we define a simple benchmarking function which we use to measure the runtime performance and memory utilization of our token generator in different scenarios.\npython\n\n```\nimport time, statistics\n\n\ndef benchmark(func, num_runs=10):\n # Warmup\n func()\n torch.cuda.synchronize()\n \n runtimes = []\n \n for _ in range(num_runs):\n # reset memory stats before each run\n torch.cuda.empty_cache()\n torch.cuda.reset_peak_memory_stats()\n torch.cuda.synchronize()\n \n start = time.perf_counter()\n _ = func()\n torch.cuda.synchronize()\n end = time.perf_counter()\n \n runtimes.append(end - start)\n \n # Get memory allocator stats from last run\n mem_stats = torch.cuda.memory_stats()\n allocated_peak = mem_stats.get('allocated_bytes.all.peak', 0)\n reserved_peak = mem_stats.get('reserved_bytes.all.peak', 0)\n f_peak = reserved_peak - allocated_peak\n f_pct = (\n 100 * f_peak / reserved_peak\n if reserved_peak > 0 else 0\n )\n \n print(f\"\\n{'='*60}\")\n print(f\"Runtime Results:\")\n print(f\" Mean: {statistics.mean(runtimes):.4f}s\")\n print(f\" Std: {statistics.stdev(runtimes):.4f}s\")\n print(f\" Min: {min(runtimes):.4f}s\")\n print(f\" Max: {max(runtimes):.4f}s\")\n\n print(f\"\\nMemory Stats:\")\n print(f\" Allocated bytes (peak): {allocated_peak / 1e9:.3f} GB\")\n print(f\" Reserved bytes (peak): {reserved_peak / 1e9:.3f} GB\")\n print(f\" Fragmentation (peak): {f_peak / 1e9:.3f} GB ({f_pct:.1f}%)\")\n print(f\"{'='*60}\\n\")\n\n\nbatch_size = 32\nfor max_seqlen in [100, 200, 400]:\n print(\n f\"Benchmarking generation with batch size {batch_size} \"\n f\"and max sequence length {max_seqlen}...\"\n )\n benchmark(\n lambda: generate_sequence(\n model, max_seqlen=max_seqlen, batch_size=batch_size\n )\n )\n```\n\nIn the table below we capture the results for a batch size of 32 and several different sequence lengths:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/2-1.webp)\nBaseline Results (By Author)\nAs the sequence length doubles, the runtime quadruples — appearing to follow a classic _O(N²)_ scaling pattern. Additionally, high memory fragmentation points to severe strain on the CUDA memory allocator, which can result in frequent memory faults and degrade runtime performance. The fragmentation results from each step asking for slightly larger tensor allocations, a pattern which ends up leaving multiple pockets of unusable memory.\nOur first optimization, KV caching, addresses the runtime complexity of our decoder model.\n### KV Caching\nOur naive generator is extremely inefficient — rather than storing and reusing the intermediate tensors from previous tokens, it recalculates the entire sequence at every step.\nWe address the computation inefficiency by using KV caching: We store and reuse the intermediate Key and Value tensors for previous tokens. KV caching reduces the runtime complexity of token generation from _O(N²)_ to _O(N)_.\nIn the following code block, we utilize the transformers library’s built-in support for KV caching to reprogram our token generation function to compute a single batch of tokens in each step.\npython\n\n```\n@torch.inference_mode()\ndef generate_sequence(model, max_seqlen, batch_size, use_cache=False):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n\n # past_key_values is used to store the cached key/values for each layer\n past_key_values = None\n\n for i in range(max_seqlen):\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n use_cache=use_cache\n )\n # update cache for next step\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n # append new token to sequence\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nThe resulting performance numbers are captured in the following table:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/3.webp)\nToken Generation With KV Caching (By Author)\nThe performance improvement is profound and, as expected, increases as a function of the sequence length.\nAlthough somewhat better than in our baseline experiment, the degree of memory fragmentation remains a concern. To address this we explore two methods, expandable memory allocations and static KV caching.\n### Expandable CUDA Memory Allocations\nTo reduce CUDA memory fragmentation, we program PyTorch to use [expandable memory segments](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf). As of the time of this writing, this memory optimization is an experimental feature and should be used with caution. Please see [the PyTorch documentation](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf) for details. To use the feature we set the following environment variable:\nbash\n\n```\nexport PYTORCH_ALLOC_CONF=\"expandable_segments:True\"\n```\n\nRerunning our benchmark results in the following table:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/4.webp)\nKV Caching With Expandable Memory Segments (By Author)\nNot only do we see a marked improvement in fragmentation, but we also get an additional (marginal) improvement in runtime performance.\n### KV Caching With StaticCache\nThe [default cache](https://huggingface.co/docs/transformers/en/kv_cache#default-cache) in HuggingFace is dynamic — it grows as the number of keys and values increases during the generation progresses. HuggingFace supports a fixed-size cache, [StaticCache](https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils#transformers.StaticCache), which pre-allocates a maximum cache size for the KV pairs and reduces strain on the CUDA memory allocator. The disadvantage of using [StaticCache](https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils#transformers.StaticCache) is that the full length of the cache participates in the attention computation at each token generation step, where irrelevant tokens are masked out. This results in a waste of computation that grows with the sequence length. For example, when generating a sequence of 400 tokens, the attention computation for each token will be run on full _400_ X _400_ -sized tensors.\nIn the code block below we enhance our sequence generator to support the use of a [StaticCache](https://huggingface.co/docs/transformers/v5.1.0/en/internal/generation_utils#transformers.StaticCache):\npython\n\n```\nche:\n\nfrom transformers import StaticCache\n\n@torch.inference_mode()\ndef generate_sequence(\n model, max_seqlen, batch_size, use_cache=False, use_static_cache=False\n):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n \n # Initialize static cache if requested\n if use_cache and use_static_cache:\n past_key_values = StaticCache(\n config=config,\n max_batch_size=batch_size,\n max_cache_len=max_seqlen,\n device=DEVICE,\n dtype=model.dtype\n )\n else:\n past_key_values = None\n \n # Initialize cache position tracking for static cache\n cache_positions = torch.arange(max_seqlen, device=DEVICE)\n \n for i in range(max_seqlen):\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n cache_position = (\n cache_positions[i:i+1] if use_static_cache else None\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n cache_position=cache_position,\n use_cache=use_cache\n )\n # update cache for next step\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n # append new token to sequence\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nThe updated results are captured below:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/5-1.webp)\nToken Generation With Static KV Cache (By Author)\nUsing a fixed-sized cache greatly improves memory utilization as indicated by the decrease in memory fragmentation. However, its impact on runtime performance is mixed — for 100 tokens it reduces performance compared to a dynamic cache, whereas for 200 and 400 tokens it boosts performance by 9% and 10%, respectively.\nThere are more advanced methods of implementing attention that optimize for memory utilization without the cost of wasted computation. In a previous post, [Optimizing Transformer Models for Variable-Length Input Sequences](https://towardsdatascience.com/optimizing-transformer-models-for-variable-length-input-sequences-19fb88fddf71/), we covered some PyTorch techniques for computing attention _sparsely_ to reduce computation waste. For production settings, libraries such as vLLM use [PagedAttention](https://blog.vllm.ai/2023/06/20/vllm.html) for maximizing memory utilization. These methods are outside the scope of this post.\nFor more details on caching in HuggingFace, please see the [caching strategies overview](https://huggingface.co/docs/transformers/en/kv_cache#cache-strategies).\n### Model Compilation\nOne of the [documented](https://huggingface.co/docs/transformers/en/kv_cache#fixed-size-cache) advantages of using a fixed-sized cache is that it allows for taking advantage of many just-in-time (JIT) optimizations.\nIn the following code block we apply our benchmark to a [PyTorch-compiled](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) version of our decoder model:\npython\n\n```\nbatch_size = 32\nmax_seqlen = 100\n\nmodel = torch.compile(model)\n\nbenchmark(\n lambda: generate_sequence(\n model,\n max_seqlen=max_seqlen,\n batch_size=batch_size,\n use_cache=True,\n use_static_cache=True\n )\n)\n```\n\nModel compilation results in an additional boost to runtime performance as shown in the table below:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/6-1.webp)\nToken Generation With torch.compile (By Author)\nNote that we can apply model compilation when using dynamic caching, as well. However, torch.compile provides the best results when the computation graph is composed of fixed-sized tensors (e.g., see [here](https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/) for more details).\n### The Performance Penalty of Early Stopping\nAn integral part of common token generators is checking for the end-of-sequence (EOS) at the end of each step. Without this test, token generators would always run for _max_seqlen,_ even if all the sequences in the batch have ended. This could result in considerable computation waste and unnecessary latency — especially when common sequence lengths are much shorter than the maximum length. In the case of our toy experiment, we wait for all the sequences in the batch to end and discontinue token generation. Production-grade implementations will commonly perform _continuous batching_ — replacing completed sequences with new prompts on the input queue.\npython\n\n```\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n \n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n```\n\nImportantly, the `.item()` call on the `stop_gpu` tensor, triggers a blocking host-device synchronization event. More specifically, in order to evaluate the conditional _if_ statement, the CPU must wait for the GPU to complete its computation and copy the contents of the tensor to host memory. While the CPU waits, it is blocked from executing the next step of the token generation loop, or more accurately, it is blocked from loading the next computation kernels onto the GPU.\nTo measure the impact of the stopping condition on runtime performance, we add instrumentation for performance profiling with [NVIDIA Nsight™ Systems (nsys)](https://docs.nvidia.com/nsight-systems/index.html) using the [torch.cuda.profiler](https://github.com/pytorch/pytorch/blob/main/torch/cuda/profiler.py) and nvtx (v0.2.14) APIs. (See our [recent post](https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/) for more details on performance profiling with nsys).\npython\n\n```\nore details on performance profiling with nsys).\n\nimport nvtx\nfrom torch.cuda import profiler\n\n@torch.inference_mode()\ndef generate_sequence(\n model, max_seqlen, batch_size, use_cache=False, use_static_cache=False\n):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n \n # Initialize static cache if requested\n if use_cache and use_static_cache:\n past_key_values = StaticCache(\n config=config,\n max_batch_size=batch_size,\n max_cache_len=max_seqlen,\n device=DEVICE,\n dtype=model.dtype\n )\n else:\n past_key_values = None\n \n # Initialize cache position tracking for static cache\n cache_positions = torch.arange(max_seqlen, device=DEVICE)\n \n for i in range(max_seqlen):\n if i == 30:\n # start nsys profiler\n torch.cuda.synchronize()\n profiler.start()\n elif i == 50:\n # stop nsys profiler\n torch.cuda.synchronize()\n profiler.stop()\n with nvtx.annotate(f\"Step {i+1}\", color=\"blue\"):\n with nvtx.annotate(\"Model Forward\", color=\"green\"):\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n cache_position = (\n cache_positions[i:i+1] if use_static_cache else None\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n cache_position=cache_position,\n use_cache=use_cache\n )\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n with nvtx.annotate(\"Check Stop Condition\", color=\"red\"):\n # checking stop condition\n if stop_gpu.item():\n print(f\"All sequences finished at step {i+1}\")\n break\n \n return all_tokens\n```\n\nWe run our script using the _cudaProfilerApi_ option to start and stop the profiler programmatically. Please see the official documentation for full details on [profiling from the nsys CLI](https://docs.nvidia.com/nsight-systems/UserGuide/index.html#profiling-from-the-cli).\nbash\n\n```\nnsys profile \\\n --capture-range=cudaProfilerApi \\\n --trace=cuda,nvtx,osrt \\\n --output=baseline \\\n python train.py\n```\n\nThe following trace, captured for a batch size of 16 and sequence length of 100, shows the GPU idling for about 110 microseconds in between steps — an eternity in the context of high-performance GPU workloads. This is a direct result of the synchronization event triggered by the EOS test.\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/7-1-1024x402.webp)\nGPU Utilization Drops Between Each Step (By Author)\nIn production-grade implementations such synchronization issues are avoided by some combination of 1) use of lower level (e.g., C/C++) code that avoids the limitation of the Python interpreter, 2) using CUDA graphs to reduce overhead of kernel loading, 3) moving conditional checks onto the GPU using [conditional nodes](https://developer.nvidia.com/blog/dynamic-control-flow-in-cuda-graphs-with-conditional-nodes/), and 4) continuously and asynchronously preparing subsequent requests while the EOS check is in progress.\nIn the next section, we demonstrate a technique for hiding the overhead of the host-device synchronization in PyTorch using [CUDA streams](https://docs.pytorch.org/docs/stable/generated/torch.cuda.Stream_class.html).\n## A CUDA Stream Optimization\nA CUDA stream is a linear sequence of operations (kernels, memory copies, etc.) that execute in order on the GPU. While operations _within_ a single stream are guaranteed to execute sequentially, operations in different streams can execute concurrently or overlap.\nIn previous posts (e.g., [here](https://towardsdatascience.com/pipelining-ai-ml-training-workloads-with-cuda-streams/) and [here](https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/)) we demonstrated the use of CUDA streams in pipelining common AI/ML workloads, e.g., executing a model on batch _N_ while preparing batch _N+1_. In this post we will use CUDA streams to enable the CPU to load the GPU kernels of step _N+1_ before checking the stopping criteria of step _N_. Contrary to our previous demonstrations of CUDA streams, our current example will not necessarily involve concurrent GPU kernel execution. \nWe implement an alternative token generation function that interleaves two CUDA streams, running the following operations iteratively:\nProgram stream _i%2_ to: **(A)** wait for stream (_i-1)%2_ to complete its generation of token _i-1,__**(**_**B)** use the updated tensors to calculate the token _i_ , **(C)** run the EOS test for token _i_ on the GPU, and **(D)** perform a (non-blocking) copy of the EOS test result to pinned memory on the CPU.\nOn the default CUDA stream, wait for stream (_i-1)%2_ to complete its generation of token _i-1._\nOn the default CUDA stream, check if the stopping criteria for token _i-1_ were met. If so, halt the generator and return. Otherwise, increment _i_ and return to step 1.\nWhereas previously, the initialization of token _i_ generation was blocked by the EOS test on token _i-1_ , the use of CUDA streams allows us to program the generation of token _i_ **before** wecheck the result of the EOS test on token _i-1_. In practice, the EOS test for token _i-1_ on the CPU runs while the GPU is computing token _i_.\npython\n\n```\n@torch.inference_mode()\ndef generate_sequence_pipelined(\n model,\n max_seqlen,\n batch_size,\n use_cache=False,\n use_static_cache=False\n):\n # Initialize prompts with BOS token\n all_tokens = torch.full(\n (batch_size, 1),\n config.bos_token_id,\n device=DEVICE,\n dtype=torch.long\n )\n finished = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool)\n past_key_values = None\n \n # Initialize static cache if requested\n if use_cache and use_static_cache:\n past_key_values = StaticCache(\n config=config,\n max_batch_size=batch_size,\n max_cache_len=max_seqlen,\n device=DEVICE,\n dtype=model.dtype\n )\n \n # Initialize cache position tracking for static cache\n cache_positions = torch.arange(max_seqlen, device=DEVICE)\n \n # Dual streams for pipelining\n streams = [torch.cuda.Stream(), torch.cuda.Stream()]\n stop_host = [\n torch.tensor(False, pin_memory=True),\n torch.tensor(False, pin_memory=True)\n ]\n \n for i in range(max_seqlen):\n curr_idx, prev_idx = i % 2, (i+1) % 2\n curr_s, prev_s = streams[curr_idx], streams[prev_idx]\n \n # Launch iteration i in current stream\n with torch.cuda.stream(curr_s):\n # program stream to wait for previous stream to complete\n curr_s.wait_stream(prev_s)\n current_input = (\n all_tokens if past_key_values is None\n else all_tokens[:, -1:]\n )\n cache_position = (\n cache_positions[i:i+1] if use_static_cache else None\n )\n outputs = model(\n current_input,\n past_key_values=past_key_values,\n cache_position=cache_position,\n use_cache=use_cache\n )\n past_key_values = outputs.past_key_values\n logits = outputs.logits[:, -1, :]\n new_tokens = torch.argmax(logits, dim=-1)\n all_tokens = torch.cat(\n [all_tokens, new_tokens.unsqueeze(-1)],\n dim=-1\n )\n \n finished |= (new_tokens == config.eos_token_id)\n stop_gpu = torch.all(finished)\n stop_host[curr_idx].copy_(stop_gpu, non_blocking=True)\n \n # Check previous iteration's stop signal\n torch.cuda.current_stream().wait_stream(prev_s)\n if stop_host[prev_idx].item():\n print(f\"All sequences finished at step {i}\")\n break\n \n return all_tokens\n```\n\nThe image below captures the nsys trace for our new token generator:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/8-1-1024x512.webp)\nConstant GPU Activity When Applying CUDA Streams (By Author)\nIn the CUDA section of the trace we can see the use of two CUDA streams, with token generation being passed back and forth in a sort of ping-pong effect: One stream generates all of the odd tokens and second all of the even tokens. The CPU is about half a step ahead of the GPU — allowing it to program step _i_ while the GPU is computing step _i-1_. The CPU-side EOS stop-check of step _i-1_(in red) occurs after step _i_ is fully programmed (and has started running).Most importantly, we now find the GPU utilization to be consistent — the idling we saw before is gone.\nThe CUDA stream interleaving results in an additional performance boost, as shown in the table below:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/9-1.webp)\nToken Generation With CUDA Streams (By Author)\nWe would expect the benefit of the ping-pong solution we have implemented to be impacted by the ratio between the GPU idle time (i.e., the overhead of kernel loading) and the kernel computation time. To test this, we fix the sequence length at 100 and rerun the benchmark for a number of batch sizes:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/10-1.webp)\nImpact of Pipelining for Varying Batch Size (By Author)\nAs expected, the highest performance gain, 11.6%, occurs when the batch size is smallest and the kernel computation load is at its lowest. As the kernel compute increases, the ratio of kernel loading to kernel compute time decreases as does the impact of CUDA stream interleaving.\nNote that there is some overhead to the use of CUDA streams. This can be demonstrated by comparing our interleaving solution to a token generator that skips the EOS test altogether:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/11-1.webp)\nOverhead of CUDA Stream Interleaving (By Author)\n### The Potential Performance Pitfalls of Using CUDA Streams\nCUDA streams should be used with extreme caution. When using the [default stream](https://docs.pytorch.org/docs/stable/generated/torch.cuda.default_stream.html) we can rely on PyTorch to perform any necessary synchronization when data is moved around. However, when using CUDA streams, we must ensure appropriate synchronization explicitly. In particular, we must ensure appropriate data transfer between the streams. Otherwise, we may experience CUDA errors (e.g., “device-side assert triggered”) — if we are lucky. If we are less lucky, we may experience data corruption without even knowing it. See the PyTorch [CUDA stream](https://docs.pytorch.org/docs/stable/notes/cuda.html#cuda-streams) documentation for more details on appropriate use.\nFor AI/ML workloads with large CUDA memory utilization, such as LLMs, another consideration is memory utilization. The PyTorch caching allocator manages memory on a per-stream basis; using multiple streams can lead to increased memory reservation and fragmentation. These could result in increased memory faults that might overshadow the potential gains from the use of streams.\n## Results\nIn the table below we summarize the runtime results of applying static caching, compilation, and pipelining on a batch of 32 sequences and a maximum sequence length of 100. The results are sorted in increasing order of performance:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/02/12-1.webp)\nToken Generation Optimization Results (By Author)\nIn the case of our toy GPT-2 model, the best results — nearly 5 times the baseline performance — are achieved when employing PyTorch compilation and the CUDA stream interleaving method discussed in this post. However, as we have seen, the impact of CUDA interleaving could vary greatly based on the properties of the workload and runtime environment, particularly on the ratio between the kernel loading time and the kernel compute time. Please be sure to run your own benchmarks before adopting this method.\n## Summary\nIn high-performance AI engineering, any hint of GPU under-utilization presents an opportunity for optimization. One of the primary optimization tools on NVIDIA GPUs is CUDA streams. In this post, we demonstrated their use in solving the idle GPU time that results from the host-device synchronization associated with early-stopping in PyTorch-native autoregressive token generation. By interleaving CUDA streams in a “ping-pong” pattern, we successfully hid the latency imposed by the EOS-check which resulted in a meaningful increase the workload’s throughput. By combining this technique with the well-known methods of model compilation and static caching, we can maximize the performance of PyTorch-native inference.\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](https://assets.insightmediagroup.io/media/1787657094227_7ymyo0.jpeg)\nWritten By\n## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Cuda](https://towardsdatascience.com/tag/cuda/)[Deep Learning](https://towardsdatascience.com/tag/deep-learning/)[Optimization](https://towardsdatascience.com/tag/optimization/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F&title=Optimizing%20Token%20Generation%20in%20PyTorch%20Decoder%20Models)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F&title=Optimizing%20Token%20Generation%20in%20PyTorch%20Decoder%20Models)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-token-generation-in-pytorch-decoder-models%2F&text=Optimizing%20Token%20Generation%20in%20PyTorch%20Decoder%20Models)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2020/07/1WFxpQU2FsvOgJKY4-BFElw.jpeg)](https://towardsdatascience.com/my-story-of-taking-the-tensorflow-developer-certification-exam-85d4fb777cbb/)\n## [My Story of Taking the TensorFlow Developer Certification Exam](https://towardsdatascience.com/my-story-of-taking-the-tensorflow-developer-certification-exam-85d4fb777cbb/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nMy overall experience of taking the exam, how I prepared for it, and what I would've done differently if I had to take the exam again.\n[Ruben Winastwan](https://towardsdatascience.com/author/marcellusruben/)\nJuly 17, 20209 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2024/04/1CLgP5Co0HELIaKO3zvGvLA.gif)](https://towardsdatascience.com/towards-infinite-llm-context-windows-e099225abaaf/)\n## [Towards infinite LLM context windows](https://towardsdatascience.com/towards-infinite-llm-context-windows-e099225abaaf/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nFrom 512 to 1M+ tokens in 5 years - LLMs have rapidly expanded their context windows. Where's the limit?\n[Krzysztof K. Zdeb](https://towardsdatascience.com/author/krzysztof.kornel/)\nApril 28, 202410 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2020/11/1akHYKaEdaOUazZBmdoNlzg.png)](https://towardsdatascience.com/resnets-residual-blocks-deep-residual-learning-a231a0ee73d2/)\n## [ResNets - Residual Blocks & Deep Residual Learning](https://towardsdatascience.com/resnets-residual-blocks-deep-residual-learning-a231a0ee73d2/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nDeep Learning harnesses the power of Big Data by building deep neural architectures that try to approximate a function f(x) that can map...\n[Prem Oommen](https://towardsdatascience.com/author/thedataoracle/)\nNovember 28, 20204 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2020/11/1rspcGSBDsRZOSXYEs6ojbg.png)](https://towardsdatascience.com/10-examples-of-awful-data-that-i-had-to-work-with-as-a-data-scientist-10894cd35828/)\n## [10 Examples of Awful Data that I had to work with as a Data Scientist](https://towardsdatascience.com/10-examples-of-awful-data-that-i-had-to-work-with-as-a-data-scientist-10894cd35828/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nA glimpse into the frustrating dark side of working with data\n[Terence Shin, MSc, MBA](https://towardsdatascience.com/author/terenceshin/)\nNovember 11, 20206 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2018/11/1Z0iOyiGUfS6nL8F2AsmByQ-scaled.jpeg)](https://towardsdatascience.com/outlier-aware-clustering-beyond-k-means-76f7bf8b4899/)\n## [Outlier-Aware Clustering: Beyond K-Means](https://towardsdatascience.com/outlier-aware-clustering-beyond-k-means-76f7bf8b4899/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nCan data science help you choose your next job? First of all, that's a terrible idea.\n[Mathieu Lemay](https://towardsdatascience.com/author/lsci/)\nNovember 29, 20182 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2020/08/0hT6KkxFUfPbEdC-d-scaled.jpg)](https://towardsdatascience.com/will-the-covid-19-vaccine-alter-my-dna-and-hook-me-to-an-ai-f7d086b94e07/)\n## [Will the COVID-19 vaccine alter my DNA and hook me to an AI?](https://towardsdatascience.com/will-the-covid-19-vaccine-alter-my-dna-and-hook-me-to-an-ai-f7d086b94e07/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nChecking your facts in the time of the pandemic\n[Alexiei Dingli](https://towardsdatascience.com/author/alexieidingli/)\nAugust 15, 20205 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/09/1onZIiGguLfbUYs3aTtmijg.jpeg)](https://towardsdatascience.com/how-to-calculate-the-mean-and-standard-deviation-normalizing-datasets-in-pytorch-704bd7d05f4c/)\n## [How to calculate the mean and standard deviation - Normalizing datasets in Pytorch](https://towardsdatascience.com/how-to-calculate-the-mean-and-standard-deviation-normalizing-datasets-in-pytorch-704bd7d05f4c/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nNeural networks converge much faster if the input data is normalized. Learn the reason why and how to implement this.\n[Jorrit Willaert](https://towardsdatascience.com/author/jorritwillaert98/)\nSeptember 24, 20213 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/01/0zwmsT2myTAXp6pxL-scaled.jpeg)](https://towardsdatascience.com/10-examples-to-master-python-dictionary-comprehensions-7aaa536f5960/)\n## [10 Examples to Master Python Dictionary Comprehensions](https://towardsdatascience.com/10-examples-to-master-python-dictionary-comprehensions-7aaa536f5960/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nA comprehensive practical guide\n[Soner Yıldırım](https://towardsdatascience.com/author/sonery/)\nJanuary 2, 20215 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/04/anya-osintsova-AkOUkGmtdc-unsplash_1.png)](https://towardsdatascience.com/an-existential-crisis-of-a-veteran-researcher-in-the-age-of-generative-ai/)\n## [An Existential Crisis of a Veteran Researcher in the Age of Generative AI](https://towardsdatascience.com/an-existential-crisis-of-a-veteran-researcher-in-the-age-of-generative-ai/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nIf you have had an existential moment in your career due to AI, let me know.\n[Pedram Ataee](https://towardsdatascience.com/author/pedram-ataee/)\nApril 23, 20255 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "body_base_path": "attachments/081c2758d42a/pages/", "body_markdown_path": "attachments/081c2758d42a/pages/archived-body.md", "local_html": "attachments/081c2758d42a/pages/archived-body.html", "body_source": "prior_capture" }