[PR #244] [CLOSED] Fix SM90 ArchSpec bug and add DG_MINIMIZE_NUM_SMS env var #308

New issue

Closed

opened 2026-10-26 21:16:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:16:06 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/244
Author: @yurekami
Created: 12/25/2025
Status: ❌ Closed

Base: main ← Head: fix/sm90-archspec-and-env-var

📝 Commits (2)

50d2eec fix: use SM90ArchSpec instead of SM100ArchSpec in sm90_bf16_k_grouped_gemm
2580e9a feat: add DG_MINIMIZE_NUM_SMS env var to control min SMS behavior

📊 Changes

4 files changed (+27 additions, -8 deletions)

View changed files

📝 csrc/jit_kernels/heuristics/sm100.hpp (+3 -1)
📝 csrc/jit_kernels/heuristics/sm90.hpp (+3 -1)
📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+4 -4)
📝 csrc/utils/system.hpp (+17 -2)

📄 Description

Summary

This PR addresses two issues:

1. Fix for Issue #242 - SM100ArchSpec Bug in SM90 Code

Problem: In sm90_bf16_k_grouped_gemm function, SM100ArchSpec was incorrectly used instead of SM90ArchSpec for TMA descriptor creation (copy-paste error)
Fix: Changed 4 occurrences in csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (lines 261, 266, 271, 272)

2. Feature for Issue #239 - Add DG_MINIMIZE_NUM_SMS Environment Variable

Problem: The min SMS optimization causes multiple kernel compilations with different SM counts, leading to unstable TTFT (time to first token)
Solution: Added DG_MINIMIZE_NUM_SMS environment variable to control this behavior
Changes:
- Added bool type support to get_env() in csrc/utils/system.hpp
- Modified should_minimize_num_sms() in sm90.hpp and sm100.hpp to check the env var

Usage:

# Disable min SMS optimization (faster startup, may use more L2 cache)
export DG_MINIMIZE_NUM_SMS=0

# Keep default behavior (default, backward compatible)
export DG_MINIMIZE_NUM_SMS=1  # or leave unset

Test plan

Verify SM90 BF16 k-grouped GEMM works correctly after the ArchSpec fix
Test with DG_MINIMIZE_NUM_SMS=0 to confirm env var disables min SMS optimization
Test with DG_MINIMIZE_NUM_SMS=1 or unset to confirm default behavior is preserved

🤖 Generated with Claude Code

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/244 **Author:** [@yurekami](https://github.com/yurekami) **Created:** 12/25/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/sm90-archspec-and-env-var` --- ### 📝 Commits (2) - [`50d2eec`](https://github.com/deepseek-ai/DeepGEMM/commit/50d2eec4c23323e3aa0acd0bd5442da524ebe3b7) fix: use SM90ArchSpec instead of SM100ArchSpec in sm90_bf16_k_grouped_gemm - [`2580e9a`](https://github.com/deepseek-ai/DeepGEMM/commit/2580e9a5d3639301d516c0de72e87c28ad347ffe) feat: add DG_MINIMIZE_NUM_SMS env var to control min SMS behavior ### 📊 Changes **4 files changed** (+27 additions, -8 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/sm100.hpp` (+3 -1) 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+3 -1) 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+4 -4) 📝 `csrc/utils/system.hpp` (+17 -2) </details> ### 📄 Description ## Summary This PR addresses two issues: ### 1. Fix for Issue #242 - SM100ArchSpec Bug in SM90 Code - **Problem:** In `sm90_bf16_k_grouped_gemm` function, `SM100ArchSpec` was incorrectly used instead of `SM90ArchSpec` for TMA descriptor creation (copy-paste error) - **Fix:** Changed 4 occurrences in `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (lines 261, 266, 271, 272) ### 2. Feature for Issue #239 - Add DG_MINIMIZE_NUM_SMS Environment Variable - **Problem:** The min SMS optimization causes multiple kernel compilations with different SM counts, leading to unstable TTFT (time to first token) - **Solution:** Added `DG_MINIMIZE_NUM_SMS` environment variable to control this behavior - **Changes:** - Added `bool` type support to `get_env()` in `csrc/utils/system.hpp` - Modified `should_minimize_num_sms()` in `sm90.hpp` and `sm100.hpp` to check the env var **Usage:** ```bash # Disable min SMS optimization (faster startup, may use more L2 cache) export DG_MINIMIZE_NUM_SMS=0 # Keep default behavior (default, backward compatible) export DG_MINIMIZE_NUM_SMS=1 # or leave unset ``` ## Test plan - [ ] Verify SM90 BF16 k-grouped GEMM works correctly after the ArchSpec fix - [ ] Test with `DG_MINIMIZE_NUM_SMS=0` to confirm env var disables min SMS optimization - [ ] Test with `DG_MINIMIZE_NUM_SMS=1` or unset to confirm default behavior is preserved 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:16:06 +00:00