[onnxruntime] ONNX Runtime WebGPU: NVIDIA GPU에서 Subgroup Matrix MatMul 성능 향상을 위한 Shared Memory Staging 최적화

2026년 10월 29일수정: 2026년 10월 29일

PR 링크: microsoft/onnxruntime#33639 상태: Merged | 변경: +2317 / -1581

들어가며

최근 ONNX Runtime의 WebGPU 백엔드에서 Subgroup Matrix MatMul 연산의 성능을 획기적으로 개선하는 Pull Request(PR)가 머지되었습니다. 이 PR은 특히 NVIDIA GPU 환경에서 행렬 곱셈(MatMul)의 속도를 높이는 데 중점을 두고 있으며, 기존의 전역 메모리(global memory) 접근 방식에서 공유 메모리(shared memory)를 활용하는 새로운 데이터 로딩 및 스테이징(staging) 전략을 도입했습니다. 본 글에서는 이 PR이 해결하고자 하는 문제, 적용된 기술적 개선 사항, 그리고 이러한 최적화가 왜 효과적인지에 대해 코드 변경 사항을 중심으로 상세히 분석해 보겠습니다.

기존 Subgroup Matrix MatMul 커널은 각 Subgroup이 A와 B 행렬을 전역 메모리에서 직접 로드했습니다. 이는 특히 대규모 행렬 연산에서 메모리 대역폭 병목 현상을 유발할 수 있었습니다. 이 PR은 이러한 병목 현상을 완화하기 위해, 워크그룹(workgroup) 내의 여러 Subgroup이 협력하여 A와 B 행렬의 일부(K 슬라이스)를 공유 메모리로 미리 로드하고, 각 Subgroup은 이 공유 메모리에서 데이터를 읽어오도록 구조를 변경했습니다. 이를 통해 전역 메모리 접근 횟수를 줄이고, GPU의 고대역폭 공유 메모리를 효율적으로 활용하여 MatMul 연산의 성능을 향상시키는 것을 목표로 합니다.

코드 분석

이번 PR의 핵심 변경 사항은 onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc 파일과 관련 WGSL 템플릿 파일에 집중되어 있습니다. 주요 변경점을 파일별로 살펴보겠습니다.

`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc`

이 파일은 Subgroup Matrix MatMul 연산을 위한 커널 생성 및 실행 로직을 담당합니다. 주요 변경 사항은 다음과 같습니다.

SubgroupMatrixMatMulProgram 생성자 변경: 새로운 파라미터 sg_grid_m, sg_grid_n, tile_k가 추가되었습니다. 이는 스테이징된 타일링(staged tiling)을 지원하기 위한 설정 값으로, 워크그룹 내 Subgroup의 그리드 크기와 타일의 K 차원을 나타냅니다.

-                              uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k)
+                              uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k,
+                              uint32_t sg_grid_m, uint32_t sg_grid_n, uint32_t tile_k)
...
-        const uint32_t split_k_; 
+        const uint32_t split_k_; 
+        const uint32_t sg_grid_m_; 
+        const uint32_t sg_grid_n_; 
+        const uint32_t tile_k_; 
  };

IsStagedTilingSupported 함수 추가: 새로운 헬퍼 함수로, 스테이징된 타일링이 현재 하드웨어 제약 조건(GPU limits) 및 설정 하에서 지원 가능한지 검증합니다. 이 함수는 N 차원이 4의 배수인지, K 분할이 1인지, 타일 크기가 유효한지, 그리고 워크그룹 크기와 공유 메모리 사용량이 제한을 초과하지 않는지 등을 확인합니다.

+bool IsStagedTilingSupported(const wgpu::Limits& limits, const SubgroupMatrixConfig& config,
+                             const SubgroupMatrixTiling& tiling, uint32_t N) {
+  const uint32_t cell_m = config.M * tiling.sg_grid_m;
+  const uint32_t cell_n = config.N * tiling.sg_grid_n;
+  if (N % 4 != 0 || tiling.split_k != 1 || cell_m == 0 || cell_n == 0 || tiling.tile_k == 0 ||
+      tiling.tile_k % config.K != 0 || tiling.tile_m % cell_m != 0 || tiling.tile_n % cell_n != 0 ||
+      tiling.tile_m / cell_m > 2 || tiling.tile_n / cell_n > 2) {
+    return false;
+  }
+  const uint64_t threads = uint64_t{config.subgroupSize} * tiling.sg_grid_m * tiling.sg_grid_n;
+  // Matches the scratch size in the template where each staged row is padded by 8 elements.
+  const uint64_t staged_elems = uint64_t{tiling.tile_m} * (tiling.tile_k + 8) + uint64_t{tiling.tile_k} * (tiling.tile_n + 8);
+  const uint64_t scratch_bytes = sizeof(uint16_t) * std::max(staged_elems, uint64_t{tiling.tile_m} * tiling.tile_n);
+  return threads <= limits.maxComputeInvocationsPerWorkgroup && threads <= limits.maxComputeWorkgroupSizeX &&
+         scratch_bytes <= limits.maxComputeWorkgroupStorageSize;
+}

SubgroupMatrixMatMulImplInternal::SelectTiling 로직 변경: 타일링 선택 로직에 스테이징된 타일링 지원 여부가 추가되었습니다. tile_k > 0인 경우 IsStagedTilingSupported 함수를 호출하여 스테이징된 타일링의 유효성을 검사합니다. 또한, 스테이징된 타일링은 공유 메모리를 사용하므로 B 행렬 패딩이 필요 없도록 needs_padded_b 계산 로직이 수정되었습니다.
```
+    if (tiling->tile_k > 0) {
+      return config.K != 0 && K % config.K == 0 && IsStagedTilingSupported(context.DeviceLimits(), config, *tiling, N);
+    }
 const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2;
-    const bool needs_padded_b = N % n_align != 0;
+    const bool needs_padded_b = !staged && N % n_align != 0;
```

SubgroupMatrixMatMulImplInternal::Compile 로직 변경: 프로그램 생성 시 sg_grid_m, sg_grid_n, tile_k 파라미터가 전달됩니다. 또한, 워크그룹 크기 계산 시 sg_grid_m과 sg_grid_n이 곱해져서 반영됩니다. 입력 텐서의 ProgramTensorMetadataDependency 타입이 스테이징 시 vec4로 변경되었습니다.

-                                        sg_mat_count_m, sg_mat_count_n, split_k)
+                                        sg_mat_count_m, sg_mat_count_n, split_k,
+                                        sg_grid_m, sg_grid_n, tile_k)
 program.SetWorkgroupSize(config.subgroupSize * split_k);
+    program.SetWorkgroupSize(config.subgroupSize * split_k * sg_grid_m * sg_grid_n);
 if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) {
   program.SetSubgroupSize(config.subgroupSize);
 }
@@ -248,8 +257,8 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl {
                   static_cast<uint32_t>(config.componentType),
                   static_cast<uint32_t>(config.resultComponentType),
                   config.M, config.N, config.K, config.subgroupSize,

                                   sg_mat_count_m, sg_mat_count_n, split_k)

   .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1},

               {b_used, ProgramTensorMetadataDependency::TypeAndRank, 1}})

                                   sg_mat_count_m, sg_mat_count_n, split_k, sg_grid_m, sg_grid_n, tile_k)

   .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, staged ? 4 : 1},

               {b_used, ProgramTensorMetadataDependency::TypeAndRank, staged ? 4 : 1}})
   .AddOutput({output, ProgramTensorMetadataDependency::Rank, 1})
   .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}, {N_b}});

GenerateSubgroupMatrixShaderCode 함수 변경: WGSL 템플릿 적용 시 sg_grid_m, sg_grid_n, tile_k 파라미터가 추가로 전달됩니다. 또한, 입력 A와 B의 ShaderVariableHelper 타입이 스테이징 시 vec4로 변경되었습니다.
```
Status GenerateSubgroupMatrixShaderCode(ShaderHelper& shader, const SubgroupMatrixConfig& config,
                                      uint32_t sg_mat_count_m, uint32_t sg_mat_count_n,
```

                                   uint32_t split_k) {

                                   uint32_t split_k, uint32_t sg_grid_m, uint32_t sg_grid_n,

                                   uint32_t tile_k, const ShaderVariableHelper& input_a,

                                   const ShaderVariableHelper& input_b) {

return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul.wgsl.template",

                          WGSL_TEMPLATE_PARAMETER(sg_grid_m, sg_grid_m),

                          WGSL_TEMPLATE_PARAMETER(sg_grid_n, sg_grid_n),
                          WGSL_TEMPLATE_PARAMETER(sg_mat_count_m, sg_mat_count_m),
                          WGSL_TEMPLATE_PARAMETER(sg_mat_count_n, sg_mat_count_n),
                          WGSL_TEMPLATE_PARAMETER(sg_mat_k, config.K),
                          WGSL_TEMPLATE_PARAMETER(sg_mat_m, config.M),
                          WGSL_TEMPLATE_PARAMETER(sg_mat_n, config.N),

                          WGSL_TEMPLATE_PARAMETER(split_k, split_k));

                          WGSL_TEMPLATE_PARAMETER(split_k, split_k),

                          WGSL_TEMPLATE_PARAMETER(tile_k, tile_k),

                          WGSL_TEMPLATE_VARIABLE(input_a, input_a),

                          WGSL_TEMPLATE_VARIABLE(input_b, input_b));

} ```

MakeStagedTilingSelector 함수 추가: NVIDIA GPU에 특화된 새로운 타일링 선택 로직을 정의합니다. 이 함수는 특정 조건(행렬 크기, K 차원 제약 등)을 만족할 때 스테이징된 타일링(tile_k > 0)을 선택합니다. 특히, 대규모 문제(M >= 128이고 특정 조건 만족 시)에는 128x128 출력 타일(4x4 Subgroup 그리드, K=64)을, 그 외에는 64x64 출력 타일(4x4 Subgroup 그리드, K=64)을 사용하도록 합니다. 이 선택 로직은 IsStagedTilingSupported 함수를 통해 하드웨어 제약 조건을 만족하는지 확인 후 적용됩니다.

+SubgroupMatrixTilingSelector MakeStagedTilingSelector(const SubgroupMatrixConfig& config) {
+  return [config, fallback = MakeDefaultTilingSelector()](
+             const ComputeContext& context, uint32_t M, uint32_t N,
+             uint32_t K, uint32_t batch) -> std::optional<SubgroupMatrixTiling> {
+    if ((M >= 64 || N >= 8192) && K % config.K == 0) {
+      const auto& limits = context.DeviceLimits();
+      const bool large = M >= 128 && uint64_t{(M + 63) / 64} * ((N + 63) / 64) >= 256;
+      const SubgroupMatrixTiling tile128{config.M * 8, config.N * 8, 1, 4, 4, 64};
+      const SubgroupMatrixTiling tile64{config.M * 4, config.N * 4, 1, 4, 4, 64};
+      if (large && IsStagedTilingSupported(limits, config, tile128, N)) {
+        return tile128;
+      }
+      if (IsStagedTilingSupported(limits, config, tile64, N)) {
+        return tile64;
+      }
+    }
+    return fallback(context, M, N, K, batch);
+  };
+}

`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template` (내용 일부)

이 템플릿 파일은 실제 WGSL 셰이더 코드를 생성하는 데 사용됩니다. 리뷰어 qjia7의 요청에 따라, 버퍼 데이터 로딩 방식이 변경되었습니다.

getByOffset 사용: 직접적인 버퍼 접근 대신 ShaderHelper의 getByOffset 메소드를 사용하여 데이터를 로드하도록 수정되었습니다. 이는 코드의 안정성과 이식성을 높이는 데 기여합니다.
```
- Please use the shader helper .getByOffset instead of raw buffer accessing to load the data.
+ Please use the shader helper .getByOffset instead of raw buffer accessing to load the data.
```
(실제 diff에서는 이 부분에 대한 구체적인 코드 변경이 명시되지 않았으나, 리뷰 코멘트를 통해 해당 변경이 이루어졌음을 알 수 있습니다. 이는 셰이더 코드 생성 로직에서 input_a와 input_b를 처리하는 방식에 영향을 미칩니다.)

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 훌륭한 최적화 및 개선이라고 할 수 있습니다.

Shared Memory Staging을 통한 성능 향상:
- 기존에는 각 Subgroup이 A와 B를 전역 메모리에서 직접 로드하여 메모리 대역폭 병목이 발생했습니다. 새로운 방식은 워크그룹 내 Subgroup들이 협력하여 A와 B의 K 슬라이스를 공유 메모리로 미리 로드합니다. Subgroup은 이 공유 메모리에서 데이터를 읽어오므로, 전역 메모리 접근 횟수가 크게 줄어듭니다.
- NVIDIA GPU의 경우, 4x4 Subgroup 그리드를 사용하여 K 차원을 따라 64 깊이로 데이터를 스테이징하고, 128x128 또는 64x64 크기의 출력 타일을 계산합니다. 이는 GPU의 고대역폭 공유 메모리를 최대한 활용하여 MatMul 연산의 처리량을 높입니다.
- PR 설명에 제시된 성능 측정 결과는 이러한 개선을 명확히 보여줍니다. Llama-3.2-1B-Instruct fp16 모델의 경우 RTX 5080에서 10.6%, RTX 4070에서 25.3%의 속도 향상을 기록했습니다. 다른 모델에서도 평균적으로 4.2% ~ 11.3%의 성능 향상이 관찰되었습니다. 특히 lm_head MatMul이 FP16으로 수행되는 q4f16 모델에서도 성능 향상이 나타난 것은 이 최적화가 특정 연산에 국한되지 않고 전반적인 MatMul 성능에 긍정적인 영향을 미침을 시사합니다.
하드웨어 특화 최적화:
- 이 최적화는 NVIDIA GPU에만 적용됩니다. 이는 각 GPU 아키텍처의 특성(예: Subgroup 기능, 공유 메모리 구조, Tensor Core 활용 등)을 고려한 세밀한 튜닝이 성능 향상에 얼마나 중요한지를 보여줍니다. 다른 벤더의 GPU에서는 기존의 안정적인 로직이 그대로 유지되어 호환성을 확보했습니다.
동적 타일링 및 제약 조건 검증:
- IsStagedTilingSupported 함수와 MakeStagedTilingSelector 로직은 단순히 새로운 기법을 적용하는 것을 넘어, 실제 하드웨어의 제약 조건(워크그룹 크기, 공유 메모리 용량 등)을 동적으로 검증합니다. 이를 통해 성능 향상이 가능한 경우에만 새로운 로직을 적용하고, 그렇지 않은 경우에는 기존의 안정적인 폴백(fallback) 로직을 사용하도록 하여 안정성을 높였습니다.
- 작은 문제의 경우, 기존의 split K 타일링이 GPU를 더 잘 채울 수 있다는 점을 고려하여 스테이징된 타일링 적용을 제한하는 등, 다양한 시나리오를 고려한 세심한 설계가 돋보입니다.
코드 품질 및 유지보수성 향상:
- 리뷰어 qjia7의 피드백에 따라 셰이더 코드에서 직접적인 버퍼 접근 대신 getByOffset 메소드를 사용하도록 변경한 것은 코드의 견고성과 유지보수성을 향상시키는 좋은 사례입니다. 이는 셰이더 프로그래밍에서 일반적인 모범 사례 중 하나입니다.
일반적인 교훈:
- 하드웨어별 최적화의 중요성: 특정 하드웨어 아키텍처의 특성을 깊이 이해하고 이를 활용하는 것이 고성능 컴퓨팅에서 필수적입니다.
- 메모리 계층 구조 활용: GPU의 공유 메모리는 전역 메모리보다 훨씬 빠르므로, 이를 효과적으로 활용하는 데이터 로딩 및 스테이징 전략은 성능 향상의 핵심입니다.
- 동적 제약 조건 검증: 최적화된 로직을 적용하기 전에 실제 실행 환경의 제약 조건을 검증하는 것은 안정성과 성능을 모두 확보하는 데 중요합니다.
- 코드 리뷰의 가치: 리뷰어의 피드백을 통해 코드의 품질과 안정성을 더욱 높일 수 있습니다.

결론

이번 PR은 ONNX Runtime WebGPU 백엔드에서 NVIDIA GPU의 Subgroup Matrix MatMul 성능을 크게 향상시키는 중요한 개선을 이루었습니다. 공유 메모리 스테이징 기법의 도입, 하드웨어 특화 최적화, 그리고 동적인 제약 조건 검증 로직은 이 최적화의 핵심입니다. 이러한 기술적 노력은 대규모 언어 모델과 같은 복잡한 AI 워크로드를 GPU에서 더 빠르고 효율적으로 실행할 수 있게 하며, ONNX Runtime의 경쟁력을 강화하는 데 기여합니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

PR Analysis 의 다른글

이전글 [sglang] SGLang에서 MiniMax-M3.1 모델의 DSpark 추론 최적화: Hopper GPU 지원 확장
현재글 : [onnxruntime] ONNX Runtime WebGPU: NVIDIA GPU에서 Subgroup Matrix MatMul 성능 향상을 위한 Shared Memory Staging 최적화
다음글 [flashinfer] FlashInfer FP8 GEMM 최적화: 윈도우 세그먼트 KG1 프로그램으로 GPU 성능 극대화

[onnxruntime] ONNX Runtime WebGPU: NVIDIA GPU에서 Subgroup Matrix MatMul 성능 향상을 위한 Shared Memory Staging 최적화

들어가며

코드 분석

`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc`

`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template` (내용 일부)

왜 이게 좋은가?

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

들어가며

코드 분석

onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc

onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template (내용 일부)

왜 이게 좋은가?

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc`

`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template` (내용 일부)