{"id":36956556,"url":"https://github.com/michaelfeil/candle-index-select-cu","last_synced_at":"2026-01-13T14:01:02.184Z","repository":{"id":328009464,"uuid":"1111967392","full_name":"michaelfeil/candle-index-select-cu","owner":"michaelfeil","description":null,"archived":false,"fork":false,"pushed_at":"2025-12-10T17:23:42.000Z","size":88,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"master","last_synced_at":"2025-12-11T03:22:14.498Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Rust","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/michaelfeil.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,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-12-08T00:36:52.000Z","updated_at":"2025-12-10T17:23:46.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/michaelfeil/candle-index-select-cu","commit_stats":null,"previous_names":["michaelfeil/candle-index-select-cu"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/michaelfeil/candle-index-select-cu","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/michaelfeil%2Fcandle-index-select-cu","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/michaelfeil%2Fcandle-index-select-cu/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/michaelfeil%2Fcandle-index-select-cu/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/michaelfeil%2Fcandle-index-select-cu/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/michaelfeil","download_url":"https://codeload.github.com/michaelfeil/candle-index-select-cu/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/michaelfeil%2Fcandle-index-select-cu/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28387596,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-13T13:42:20.960Z","status":"ssl_error","status_checked_at":"2026-01-13T13:42:03.276Z","response_time":56,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"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":[],"created_at":"2026-01-13T14:00:59.234Z","updated_at":"2026-01-13T14:01:02.168Z","avatar_url":"https://github.com/michaelfeil.png","language":"Rust","funding_links":[],"categories":[],"sub_categories":[],"readme":"# candle-index-select-cu\n\nFast CUDA `index_select` for [Candle](https://github.com/huggingface/candle). Inspired by https://github.com/huggingface/candle-layer-norm project.\n\nThis crate provides a specialized CUDA kernel for `index_select` along\ndimension 0 on multi-dimensional tensors, with:\n\n- **Optimized memory access patterns** using vectorized loads (float4 for F32, half2 for F16)\n- **F32 and F16 data types** with dtype-specific optimizations\n- **CUDA-only implementation** with automatic fallback to candle's builtin\n- **Dimension 0 specialization** for row-based indexing operations\n- **Contiguous memory layout optimization** for better cache utilization\n- **Significant speedups** on large tensors (up to 4-22× vs candle, 4× vs PyTorch 2.1) \n- **Tested** against candle index_select for correctness.\n- faster only if dim=0 and tensor dim=0 cols is divisible by 2 for fp16 (half2) and 4 for fp32 (float4) memory ops.\n\n### Benchmark Results (H100 80GB)\n\n#### vs candle 0.9.1 native kernels\n\n| Input Shape            | Out Rows | DType | candle 0.9.1 | candle-index-select-cu | Speedup |\n|------------------------|----------|-------|--------------|------------------------|---------|\n| [100, 128]             | 200      | F32   | 16.436 µs    | 12.645 µs              | **1.30×** |\n| [100, 128]             | 200      | F16   | 16.676 µs    | 12.505 µs              | **1.33×** |\n| [16000, 1024]          | 12000    | F32   | 340.410 µs   | 46.250 µs              | **7.36×** |\n| [16000, 1024]          | 12000    | F16   | 106.553 µs   | 28.435 µs              | **3.75×** |\n| [16000, 1024]          | 70000    | F32   | 1.980 ms     | 204.540 µs             | **9.68×** |\n| [16000, 1024]          | 70000    | F16   | 1.073 ms     | 106.728 µs             | **10.06×** |\n| [100000, 2048]         | 500000   | F32   | 45.636 ms    | 2.828 ms               | **16.14×** |\n| [100000, 2048]         | 500000   | F16   | 33.665 ms    | 1.469 ms               | **22.92×** |\n| [10, 100, 128]         | 200      | F32   | 33.675 µs    | 16.217 µs              | **2.08×** |\n| [10, 100, 128]         | 200      | F16   | 33.953 µs    | 17.115 µs              | **1.98×** |\n| [2000, 64, 256]        | 10000    | F32   | 3.737 ms     | 432.668 µs             | **8.64×** |\n| [2000, 64, 256]        | 10000    | F16   | 2.880 ms     | 211.468 µs             | **13.62×** |\n\n*Benchmarks run on NVIDIA H100 80GB HBM3.*\n\n#### vs PyTorch 2.1.0a0+b5021ba (CUDA 12.1)\n\n| Input Shape            | Out Rows | DType | PyTorch      | candle-index-select-cu | Speedup |\n|------------------------|----------|-------|--------------|------------------------|---------|\n| [100, 128]             | 200      | F32   | 15.256 µs    | 12.645 µs              | **1.21×** |\n| [100, 128]             | 200      | F16   | 14.928 µs    | 12.505 µs              | **1.19×** |\n| [16000, 1024]          | 70000    | F32   | 515.830 µs   | 204.540 µs             | **2.52×** |\n| [16000, 1024]          | 70000    | F16   | 489.939 µs   | 106.728 µs             | **4.59×** |\n| [100000, 2048]         | 500000   | F32   | 7.958 ms     | 2.828 ms               | **2.81×** |\n| [100000, 2048]         | 500000   | F16   | 7.127 ms     | 1.469 ms               | **4.85×** |\n| [10, 100, 128]         | 200      | F32   | 25.007 µs    | 16.217 µs              | **1.54×** |\n| [10, 100, 128]         | 200      | F16   | 24.808 µs    | 17.115 µs              | **1.45×** |\n| [2000, 64, 256]        | 10000    | F32   | 1.035 ms     | 432.668 µs             | **2.39×** |\n| [2000, 64, 256]        | 10000    | F16   | 1.005 ms     | 211.468 µs             | **4.75×** |\n\n*Benchmarks depend on Pytorch version.*\n\n## Usage\n\nAdd to your `Cargo.toml`:\n\n```toml\n[dependencies]\ncandle-index-select-cu = \"0.0.1\"\ncandle-core = { version = \"0.9\", features = [\"cuda\"] }\n```\n\nThen use it in your code:\n\n```rust\nuse candle_index_select_cu::index_select;\nuse candle_core::{Device, Tensor, DType};\n\nlet device = Device::new_cuda(0)?;\nlet x = Tensor::randn(0f32, 1.0, (1000, 512), \u0026device)?;\nlet indices = Tensor::from_vec(vec![0u32, 5, 10, 15], 4, \u0026device)?;\n\n// Fast path for 2D + dim 0 + f32/f16 + u32 indices\n// Falls back to candle's builtin for other cases\nlet result = index_select(\u0026x, \u0026indices, 0)?;\n```\n\n### Recommended Integration with Feature Flags\n\nFor better compatibility to CPU and other backends, consider wrapping the usage behind a CUDA feature flag:\n\n```toml\n[dependencies]\ncandle-index-select-cu = { version = \"0.0.1\", optional = true }\ncandle-core = \"0.9\"\n\n[features]\ndefault = []\ncuda = [\"candle-core/cuda\", \"candle-index-select-cu\"]\n```\n\n```rust\n#[cfg(feature = \"cuda\")]\nuse candle_index_select_cu;\n\n#[inline]\npub fn index_select(tensor: \u0026Tensor, ids: \u0026Tensor, dim: usize) -\u003e Result\u003cTensor\u003e {\n    #[cfg(not(feature = \"cuda\"))]\n    {\n        tensor.index_select(ids, dim)\n    }\n    #[cfg(feature = \"cuda\")]\n    {\n        candle_index_select_cu::index_select(tensor, ids, dim)\n    }\n}\n```\n\n## Feature Flags\n\nThis crate supports both older and newer versions of candle/cudarc:\n\n| Feature | Candle Version | cudarc Version | Default |\n|---------|---------------|----------------|---------|\n| `cuda-12` | 0.9+ | 0.16+ | ✅ |\n| `cuda-11` | pre-0.9 (0.6, 0.7, 0.8) | pre-0.16 | |\n\n### For candle 0.9+ (default)\n\n```toml\n[dependencies]\ncandle-index-select-cu = { git = \"https://github.com/michaelfeil/candle-index-select-cu\" }\n# best to globally pin candle -\u003e via patches\n```\n\n### For older candle versions (pre-0.9)\n\n```toml\n[dependencies]\ncandle-index-select-cu = { git = \"https://github.com/michaelfeil/candle-index-select-cu\", default-features = false, features = [\"cuda-11\"] }\n```\n\n## Fast Path Conditions\n\nThe CUDA kernel is used when all of these conditions are met:\n\n- Device: CUDA\n- Input tensor: rank \u003e= 2 (2D, 3D, 4D, etc.), fully contiguous\n- Indices tensor: rank 1, contiguous, DType::U32\n- Dimension: 0\n- Input dtype: F16 or F32\n\nFor all other cases, the function falls back to candle's builtin `index_select`.\n\n## Running Tests\n\n```bash\ncargo test\n```\n\n## Running Benchmarks\n\n```bash\ncargo bench\n```\n\n## Profiling\n\n### NSys\n```\ncargo build --release --bin profile_index_select \u0026\u0026 nsys profile -t cuda,osrt --stats=true -o nsys_index_select   cargo/release/profile_index_select\n```\n\n### NCU\n```\ncargo build --release --bin profile_index_select \u0026\u0026 CUDA_LAUNCH_BLOCKING=1 ncu --set full     --devices 0     --target-processes all     --kernel-name-base demangled     -k \"regex:.*index_select.*\"     --launch-skip 1500     --launch-count 1     -o index_select_profile  cargo/release/profile_index_select\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmichaelfeil%2Fcandle-index-select-cu","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmichaelfeil%2Fcandle-index-select-cu","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmichaelfeil%2Fcandle-index-select-cu/lists"}