{"id":26187723,"url":"https://github.com/shreyansh26/accelerating-cross-encoder-inference","last_synced_at":"2026-07-02T05:31:46.970Z","repository":{"id":280305252,"uuid":"936762039","full_name":"shreyansh26/Accelerating-Cross-Encoder-Inference","owner":"shreyansh26","description":"Leveraging torch.compile to accelerate cross-encoder inference","archived":false,"fork":false,"pushed_at":"2025-03-03T04:12:23.000Z","size":4027,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-11T23:51:24.300Z","etag":null,"topics":["cross-encoder","inference-optimization","jina","mlsys","torch-compile"],"latest_commit_sha":null,"homepage":"https://shreyansh26.github.io/post/2025-03-02_cross-encoder-inference-torch-compile/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/shreyansh26.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2025-02-21T16:35:51.000Z","updated_at":"2025-03-03T04:12:27.000Z","dependencies_parsed_at":"2025-03-02T16:34:43.519Z","dependency_job_id":"555abd85-4e4e-4f9d-bf84-1a8ca98c7087","html_url":"https://github.com/shreyansh26/Accelerating-Cross-Encoder-Inference","commit_stats":null,"previous_names":["shreyansh26/accelerating-cross-encoder-inference"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/shreyansh26/Accelerating-Cross-Encoder-Inference","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/shreyansh26%2FAccelerating-Cross-Encoder-Inference","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/shreyansh26%2FAccelerating-Cross-Encoder-Inference/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/shreyansh26%2FAccelerating-Cross-Encoder-Inference/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/shreyansh26%2FAccelerating-Cross-Encoder-Inference/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/shreyansh26","download_url":"https://codeload.github.com/shreyansh26/Accelerating-Cross-Encoder-Inference/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/shreyansh26%2FAccelerating-Cross-Encoder-Inference/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":35034984,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-07-02T02:00:06.368Z","response_time":173,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cross-encoder","inference-optimization","jina","mlsys","torch-compile"],"created_at":"2025-03-11T23:51:16.857Z","updated_at":"2026-07-02T05:31:46.945Z","avatar_url":"https://github.com/shreyansh26.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Accelerating Cross Encoder inference with torch.compile\n\n## Overview\nThis project demonstrates optimizing the inference of a Cross Encoder model, namely [jinaai/jina-reranker-v2-base-multilingual](https://huggingface.co/jinaai/jina-reranker-v2-base-multilingual), by leveraging torch.compile. The scripts compare baseline performance against a torch.compile-optimized version using a custom padding approach.\n\n**Blog post describing the approach in detail** - https://shreyansh26.github.io/post/2025-03-02_cross-encoder-inference-torch-compile/\n\n## Setup\n- Python 3.8+\n- PyTorch with CUDA support\n- Sentence Transformers library\n- Model: jinaai/jina-reranker-v2-base-multilingual\n\n## Scripts\n\n### bench_basic.py\nRuns a baseline benchmark with the standard CrossEncoder and Flash Attention enabled.\n\n### bench_torch_compile.py\nFocuses on the torch.compile approach with some custom padding and torch.compile optimizations.\n\n### bench_combined.py\nCompares the baseline with torch.compile optimized version.\n\n## Implementation Details\n\n- **Batching with custom padding**: The custom `DynamicCrossEncoder` pads tokenized inputs to a bucket length (multiples of 16), to lower the number of dynamic lengths that torch.compile has to capture.\n\n- **Sorted Inputs**: Sorting the inputs before batching allows the sequences in the batch to be of similar lengths hence less padding tokens to be processed.\n\n- **torch.compile**: The model's forward function is compiled using `torch.compile` with the `inductor` backend, enabling dynamic shape handling and reducing latency.\n\n## Speedup Analysis\n\nThe torch.compile optimized version shows significant speedups compared to the baseline (batch size 64, H100 GPU):\n\n| Setup                                      | Sorted Inputs (s)       | Unsorted Inputs (s)      |\n| ------------------------------------------ | ----------------------- | ------------------------ |\n| Base (with Flash Attention)                | 0.2658 ± 0.0119         | 0.2961 ± 0.0089          |\n| torch.compile                              | 0.2089 ± 0.0196         | 0.2595 ± 0.0077          |\n\nThis reflects roughly a 20-25% reduction in inference latency under sorted inputs, with similar gains observed for unsorted inputs.\n\n## How to Run\n1. Ensure your environment has CUDA and the required libraries installed:\n   - `pip install sentence-transformers torch`\n2. Execute the benchmark scripts:\n   - `CUDA_VISIBLE_DEVICES=0 python bench_basic.py`\n   - `CUDA_VISIBLE_DEVICES=0 python bench_torch_compile.py`\n   - `CUDA_VISIBLE_DEVICES=0 python bench_combined.py`\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fshreyansh26%2Faccelerating-cross-encoder-inference","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fshreyansh26%2Faccelerating-cross-encoder-inference","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fshreyansh26%2Faccelerating-cross-encoder-inference/lists"}