[Model Bash] DeepSeek R1 BF16 Min Latency QKV A GEMM (0.5% E2E Speedup)

[robertgshaw2-redhat]

Purpose

• add min latency bf16 qkv_a_gemm
• adapted from sgl: * https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu, which adapted from trtllm
• disappointing E2E win, but sets up using PDL to overlap with AR in future PR

Test Plan

• lm eval

local-completions ({'model': 'nvidia/DeepSeek-R1-NVFP4', 'base_url': 'http://localhost:7000/v1/completions', 'num_concurrent': 1000, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: None, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9568|±  |0.0056|
|     |       |strict-match    |     5|exact_match|↑  |0.9553|±  |0.0057|

• benchmark

sweep:
	just benchmark 1 10 && \
	just benchmark 4 40 && \
	just benchmark 7 70 && \
	just benchmark 8 80 && \
	just benchmark 12 120 && \
	just benchmark 16 160

Test Result

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a highly optimized fused GEMM kernel for DeepSeek V2/V3 models on Hopper architecture GPUs. The changes include the CUDA kernel itself, build system modifications in CMake, and integration into the Python codebase. The integration logic correctly checks for all conditions required to use this specialized kernel. My main feedback is to add error handling for CUDA API calls in the new kernel file to make it more robust.

[gemini-code-assist]

The CUDA runtime API calls cudaGetDevice and cudaDeviceGetAttribute are not checked for errors. If these calls fail, the function might return an incorrect SM version (e.g., 0), which could lead to silent failures or incorrect behavior downstream (e.g., dsv3_fused_a_gemm failing with a generic "required CUDA ARCH >= SM_90" message, or optimizations being silently disabled). It is recommended to add error checking for these CUDA calls, for example by using a macro that checks the cudaError_t return value and throws an exception on failure.

[mgoin]

LGTM just need to fix sm120 restriction

[eugr]

@robertgshaw2-redhat - this commit causes vLLM to fail on startup on DGX Spark (sm121), vLLM built with TORCH_CUDA_ARCH_LIST=12.1a:

Traceback (most recent call last):
  File "/usr/local/bin/vllm", line 4, in <module>
    from vllm.entrypoints.cli.main import main
  File "/usr/local/lib/python3.12/dist-packages/vllm/entrypoints/cli/__init__.py", line 3, in <module>
    from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand
  File "/usr/local/lib/python3.12/dist-packages/vllm/entrypoints/cli/benchmark/latency.py", line 5, in <module>
    from vllm.benchmarks.latency import add_cli_args, main
  File "/usr/local/lib/python3.12/dist-packages/vllm/benchmarks/latency.py", line 16, in <module>
    from vllm.engine.arg_utils import EngineArgs
  File "/usr/local/lib/python3.12/dist-packages/vllm/engine/arg_utils.py", line 35, in <module>
    from vllm.config import (
  File "/usr/local/lib/python3.12/dist-packages/vllm/config/__init__.py", line 5, in <module>
    from vllm.config.cache import CacheConfig
  File "/usr/local/lib/python3.12/dist-packages/vllm/config/cache.py", line 13, in <module>
    from vllm.utils.mem_utils import format_gib, get_cpu_memory
  File "/usr/local/lib/python3.12/dist-packages/vllm/utils/mem_utils.py", line 14, in <module>
    from vllm.platforms import current_platform
  File "/usr/local/lib/python3.12/dist-packages/vllm/platforms/__init__.py", line 252, in __getattr__
    _current_platform = resolve_obj_by_qualname(platform_cls_qualname)()
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/vllm/utils/import_utils.py", line 111, in resolve_obj_by_qualname
    module = importlib.import_module(module_name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/importlib/__init__.py", line 90, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py", line 16, in <module>
    import vllm._C  # noqa
    ^^^^^^^^^^^^^^
ImportError: /usr/local/lib/python3.12/dist-packages/vllm/_C.abi3.so: undefined symbol: _Z17dsv3_fused_a_gemmRN2at6TensorERKS0_S3_

@mgoin, @johnnynunez - FYI.

[SurealCereal]

@robertgshaw2-redhat - this commit causes vLLM to fail on startup on DGX Spark (sm121), vLLM built with TORCH_CUDA_ARCH_LIST=12.1a:

It fails for me on startup with TORCH_CUDA_ARCH_LIST=12.0a running on an RTX PRO 6000 also. Here is the stack trace:

vllm  | Traceback (most recent call last):
vllm  |   File "/usr/local/bin/vllm", line 4, in <module>
vllm  |     from vllm.entrypoints.cli.main import main
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/entrypoints/cli/__init__.py", line 3, in <module>
vllm  |     from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/entrypoints/cli/benchmark/latency.py", line 5, in <module>
vllm  |     from vllm.benchmarks.latency import add_cli_args, main
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/benchmarks/latency.py", line 16, in <module>
vllm  |     from vllm.engine.arg_utils import EngineArgs
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/engine/arg_utils.py", line 35, in <module>
vllm  |     from vllm.config import (
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/config/__init__.py", line 5, in <module>
vllm  |     from vllm.config.cache import CacheConfig
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/config/cache.py", line 13, in <module>
vllm  |     from vllm.utils.mem_utils import format_gib, get_cpu_memory
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/mem_utils.py", line 14, in <module>
vllm  |     from vllm.platforms import current_platform
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/platforms/__init__.py", line 252, in __getattr__
vllm  |     _current_platform = resolve_obj_by_qualname(platform_cls_qualname)()
vllm  |                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/import_utils.py", line 111, in resolve_obj_by_qualname
vllm  |     module = importlib.import_module(module_name)
vllm  |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
vllm  |   File "/usr/lib/python3.12/importlib/__init__.py", line 90, in import_module
vllm  |     return _bootstrap._gcd_import(name[level:], package, level)
vllm  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
vllm  |   File "/usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py", line 16, in <module>
vllm  |     import vllm._C  # noqa
vllm  |     ^^^^^^^^^^^^^^
vllm  | ImportError: /usr/local/lib/python3.12/dist-packages/vllm/_C.abi3.so: undefined symbol: _Z17dsv3_fused_a_gemmRN2at6TensorERKS0_S3_
vllm exited with code 1 (restarting)

[mgoin]

Thank you for reporting @eugr @SurealCereal and sorry for the disruption. I should have a fix here #35123