[Feature] Add DoSA prefill indexer on SM100

.buildkite/test_areas/kernels.yaml

@@ -43,6 +43,8 @@ steps:
   - vllm/kernels/triton/
   - tests/kernels/test_concat_mla_q.py
   - tests/kernels/test_fused_qk_norm_rope_gate.py
+  - vllm/model_executor/layers/rotary_embedding/
+  - vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
   commands:
     - >-
       pytest -v -s
@@ -70,6 +72,7 @@ steps:
       - vllm/model_executor/layers/fused_moe/router/gate_linear.py
       - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
       - vllm/model_executor/layers/rotary_embedding/
+      - vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
       - tests/kernels/test_fused_qk_norm_rope_gate.py
@@ -689,6 +692,26 @@ steps:
     - pytest -v -s tests/kernels/attention/test_minimax_m3.py -k prefill_sparse_attention_correctness
     - pytest -v -s tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py
 
+- label: ":nvidia: (B200) DoSA Kernels"
+  key: kernels-dosa-test-b200
+  timeout_in_minutes: 20
+  working_dir: "/vllm-workspace/"
+  device: b200-k8s
+  num_devices: 1
+  source_file_dependencies:
+  - csrc/attention/dosa/
+  - cmake/dosa.cmake
+  - vllm/model_executor/layers/dosa.py
+  - vllm/model_executor/layers/dosa_ops.py
+  - vllm/model_executor/layers/sparse_attn_indexer.py
+  - vllm/v1/attention/backends/mla/indexer.py
+  - vllm/v1/attention/backends/mla/sparse_utils.py
+  - vllm/config/attention.py
+  - tests/kernels/attention/test_dosa.py
+  - tests/conftest.py
+  commands:
+    - pytest -v -s tests/kernels/attention/test_dosa.py
+
 - label: ":nvidia: (B200) Kernels"
   key: kernels-b200
   automatic_shard: true

CMakeLists.txt

@@ -1715,6 +1715,7 @@ endif()
 # For CUDA we also build and ship some external projects.
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
+    include(cmake/dosa.cmake)
     include(cmake/external_projects/deepselect.cmake)
     include(cmake/external_projects/fmha_sm100.cmake)
     include(cmake/external_projects/flashmla.cmake)

cmake/dosa.cmake

@@ -0,0 +1,56 @@
+# Optional AOT extension; operators share vLLM's existing _C Torch namespace.
+set(_dosa_root "${CMAKE_CURRENT_SOURCE_DIR}/csrc/attention/dosa")
+set(_dosa_sources "${_dosa_root}/module.cpp")
+cuda_archs_loose_intersection(DOSA_ARCHS "10.0a" "${CUDA_ARCHS}")
+if(DOSA_ARCHS AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
+  set(_dosa_generated "${CMAKE_CURRENT_BINARY_DIR}/dosa")
+  file(MAKE_DIRECTORY "${_dosa_generated}")
+  foreach(_heads 32 64)
+    foreach(_role compact direct_abs high low planner refine sample select)
+      set(_namespace "dosa_h${_heads}")
+      set(_source "${_dosa_root}/${_role}.cu")
+      set(_role_defines "")
+      if(_role STREQUAL "direct_abs")
+        set(_namespace "dosa_h${_heads}_direct")
+        set(_source "${_dosa_root}/planner.cu")
+        set(_role_defines "#define DOSA_DIRECT_ABS 1\n")
+      endif()
+      set(_unit "${_dosa_generated}/h${_heads}_${_role}.cu")
+      string(CONCAT _dosa_unit_content
+        "${_role_defines}"
+        "#define DOSA_SCORE_HEADS ${_heads}\n"
+        "#define DOSA_NAMESPACE ${_namespace}\n"
+        "#define DOSA_OP_PREFIX \"dosa_h${_heads}_${_role}_\"\n"
+        "#include \"${_source}\"\n")
+      file(GENERATE OUTPUT "${_unit}" CONTENT "${_dosa_unit_content}")
+      set_gencode_flags_for_srcs(SRCS "${_unit}" CUDA_ARCHS "${DOSA_ARCHS}")
+      set_source_files_properties("${_unit}" PROPERTIES INCLUDE_DIRECTORIES
+        "${deepgemm_SOURCE_DIR}/deep_gemm/include;${deepgemm_SOURCE_DIR}/third-party/cutlass/include")
+      set_property(SOURCE "${_unit}" APPEND PROPERTY COMPILE_OPTIONS
+        -O3 -lineinfo --expt-relaxed-constexpr -std=c++20
+        --prec-div=true --prec-sqrt=true --ftz=false --fmad=true)
+      if(_role MATCHES "^(low|direct_abs)$")
+        set_property(SOURCE "${_unit}" APPEND PROPERTY COMPILE_OPTIONS --fmad=false)
+      endif()
+      if(_role STREQUAL "low")
+        set_property(SOURCE "${_unit}" APPEND PROPERTY COMPILE_OPTIONS --expt-extended-lambda)
+      endif()
+      list(APPEND _dosa_sources "${_unit}")
+    endforeach()
+  endforeach()
+endif()
+
+set(_dosa_gpu_flags ${VLLM_GPU_FLAGS})
+list(REMOVE_ITEM _dosa_gpu_flags --use_fast_math -use_fast_math)
+define_extension_target(
+  _dosa_C
+  DESTINATION vllm
+  LANGUAGE CUDA
+  SOURCES ${_dosa_sources}
+  COMPILE_FLAGS ${_dosa_gpu_flags}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  USE_SABI 3
+  WITH_SOABI)
+target_compile_definitions(_dosa_C PRIVATE
+  TORCH_TARGET_VERSION=0x020B000000000000ULL
+  USE_CUDA CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)