[Perf] Use 2D-grid to eliminate divmod in W8W8 group quant

[jiahanc]

Purpose

Replace the 1D-grid (global_group_id % padded_gpr, / padded_gpr) divmod with a 2D grid + template-constant index unpack. Since the tile sizes are compile-time constants, the compiler can lower the coordinate math to simple bit operations instead of emitting runtime division/modulo setup.

Micro-benchmark on DeepSeek V4 Pro shape

K	M	baseline us	opt us	speedup
7168	8192	36.95	33.62	1.099x
7168	1024	5.68	5.26	1.080x
7168	256	2.76	2.66	1.035x
7168	64	2.10	1.98	1.061x
7168	8	1.93	1.89	1.022x
7168	1	1.85	1.84	1.005x
2048	8192	10.84	10.05	1.079x
2048	1024	2.86	2.76	1.034x
2048	256	2.11	1.99	1.058x
2048	64	1.98	1.87	1.057x
2048	8	1.86	1.85	1.005x
2048	1	1.85	1.85	1.000x
1536	8192	8.55	7.96	1.075x
1536	1024	2.65	2.54	1.043x
1536	256	2.07	1.95	1.061x
1536	64	1.94	1.87	1.040x
1536	8	1.86	1.86	1.000x
1536	1	1.83	1.83	1.000x
384	8192	4.17	3.94	1.057x
384	1024	2.21	2.09	1.058x
384	256	2.07	1.94	1.067x
384	64	2.01	1.89	1.065x
384	8	1.88	1.87	1.005x
384	1	1.84	1.84	1.000x

Test Plan

lm_eval gsm8k

Test Result

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9538|±  |0.0058|
|     |       |strict-match    |     5|exact_match|↑  |0.9538|±  |0.0058|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[jiahanc]

e2e benchmark on DeepSeek V4 Flash TP4
Concurrency 1 Baseline

============ Serving Benchmark Result ============
Successful requests:                     4         
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  36.79     
Total input tokens:                      4096      
Total generated tokens:                  4096      
Request throughput (req/s):              0.11      
Output token throughput (tok/s):         111.33    
Peak output token throughput (tok/s):    4.00      
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          222.66    
---------------Time to First Token----------------
Mean TTFT (ms):                          134.17    
Median TTFT (ms):                        133.98    
P99 TTFT (ms):                           136.45    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          8.86      
Median TPOT (ms):                        8.86      
P99 TPOT (ms):                           8.87      
==================================================

Concurrency 1 OPT

============ Serving Benchmark Result ============
Successful requests:                     4         
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  36.50     
Total input tokens:                      4096      
Total generated tokens:                  4096      
Request throughput (req/s):              0.11      
Output token throughput (tok/s):         112.22    
Peak output token throughput (tok/s):    4.00      
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          224.45    
---------------Time to First Token----------------
Mean TTFT (ms):                          134.57    
Median TTFT (ms):                        133.85    
P99 TTFT (ms):                           137.75    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          8.79      
Median TPOT (ms):                        8.79      
P99 TPOT (ms):                           8.79      
==================================================

[jiahanc]

e2e benchmark on DeepSeek V4 Flash TP4
Concurrency 1024 Baseline

============ Serving Benchmark Result ============
Successful requests:                     4096      
Failed requests:                         0         
Maximum request concurrency:             1024      
Benchmark duration (s):                  290.90    
Total input tokens:                      4194304   
Total generated tokens:                  4194304   
Request throughput (req/s):              14.08     
Output token throughput (tok/s):         14418.23  
Peak output token throughput (tok/s):    568.00    
Peak concurrent requests:                1079.00   
Total token throughput (tok/s):          28836.46  
---------------Time to First Token----------------
Mean TTFT (ms):                          3563.84   
Median TTFT (ms):                        659.25    
P99 TTFT (ms):                           21673.61  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          67.08     
Median TPOT (ms):                        68.93     
P99 TPOT (ms):                           70.56     
==================================================

Concurrency 1024 OPT

============ Serving Benchmark Result ============
Successful requests:                     4096      
Failed requests:                         0         
Maximum request concurrency:             1024      
Benchmark duration (s):                  289.16    
Total input tokens:                      4194304   
Total generated tokens:                  4194304   
Request throughput (req/s):              14.17     
Output token throughput (tok/s):         14505.22  
Peak output token throughput (tok/s):    547.00    
Peak concurrent requests:                1080.00   
Total token throughput (tok/s):          29010.44  
---------------Time to First Token----------------
Mean TTFT (ms):                          3558.47   
Median TTFT (ms):                        723.74    
P99 TTFT (ms):                           21689.62  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          66.66     
Median TPOT (ms):                        68.78     
P99 TPOT (ms):                           69.36     
==================================================

[gemini-code-assist]

Code Review

This pull request refactors the 8-bit packed per-token group quantization kernel to use a 2D grid layout and template-based block dimensions. It introduces the GetGroupsPerBlockX helper function to determine optimal tiling and updates the kernel launch logic to support varying X and Y block configurations. Feedback suggests improving the robustness of the kernel launch macro by explicitly checking for all supported tiling factors and adding an error path for unsupported values.

[gemini-code-assist]

The LAUNCH_REG_KERNEL macro assumes that kx can only be 16, 8, or 4 based on the GetGroupsPerBlockX implementation. While this is currently true because padded_groups_per_row is always a multiple of 4, it would be safer and more robust to include a static assertion or a more explicit check to ensure that kx * ry always equals the expected number of groups per block (16) for which the kernel is optimized. If kx were to ever be a value not handled by the if-else chain, it would default to the 4, 4 case which might be incorrect for the actual kx value.

#define LAUNCH_REG_KERNEL(T, DST_DTYPE)              \
  do {                                               \
    if (kx == 16) {                                  \
      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1);   \
    } else if (kx == 8) {                            \
      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2);    \
    } else if (kx == 4) {                            \
      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4);    \
    } else {                                         \
      STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
    }                                                \
  } while (0)

[mergify]

Hi @jiahanc, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10