Refactor device info retrieval into shared utility with global caching

[Copilot]

PR Category

Refactor

Type of Change

Performance Optimization | Refactor

Description

Extract device info retrieval logic from flash_attn_varlen_func.cpp and consolidate with mm.cpp into a shared, globally-cached utility in flag_gems::utils.

Changes:

• Add DeviceInfo struct and cached getter to include/flag_gems/utils.h and lib/utils.cpp
• Replace per-call cudaGetDevice/cudaGetDeviceProperties in flash_attn_varlen_func.cpp with utils::get_sm_count()
• Remove duplicate DeviceInfo implementation from mm.cpp, use shared utility

Usage:

// Before: called every invocation
int dev_id = 0;
cudaGetDevice(&dev_id);
cudaDeviceProp prop {};
cudaGetDeviceProperties(&prop, dev_id);
int num_sms = prop.multiProcessorCount;

// After: cached once per process
int num_sms = flag_gems::utils::get_sm_count();

Provides get_device_id(), get_l2_cache_size(), get_sm_count(), get_major() with thread-safe initialization (C++17 static local).

Issue

Resolves device info retrieval inefficiency - properties now queried once per process instead of every function call.

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Eliminates repeated CUDA runtime API calls for device properties. Fallback defaults (A100: 108 SMs, 40MB L2, major=8) if queries fail.

Original prompt

flash_attn_varlen_func.cpp 中获取device 信息的逻辑，当前是每次调用都会获取一遍，这个逻辑能不能抽取出来，连同mm.cpp 中的device 获取逻辑，放到一个公共的地方，并且可以是全局缓存的。

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.