[Optimization]【Hackathon 10th Spring No.49】Port ngram_match and hybrid_mtp_ngram kernels to CUDA

[cloudforge1]

Motivation

Speculative decoding in FastDeploy uses n-gram matching (ngram_match and hybrid_mtp_ngram) to propose draft tokens.
Both kernels currently run on CPU, requiring synchronous Device→CPU→Device data copies for ~10 tensors per call.
These forced CUDA stream synchronizations are a significant latency bottleneck.

This PR ports both kernels to CUDA with a two-phase parallel architecture, eliminating all device↔host data transfers and parallelizing the sliding-window ngram search across batch items and sequence positions.

Addresses Hackathon 10th Spring No.49 — "Speculative Decoding Kernel for FastDeploy".

Related RFC: community#1213

Modifications

Architecture: Two-Phase Parallel Kernel

Phase 1 — Parallel Search <<<bsz, 256>>>:

• One CUDA block per batch item, 256 threads per block
• Each thread handles a slice of the sequence via strided sliding-window ngram search
• atomicMin64 CAS loop ensures leftmost-match semantics (matching position written atomically to shared NgramMatchResult)
• Block-level reduction via __shared__ memory (s_min_pos) — threads find local candidates, block picks the leftmost

Phase 2 — Serial Gather <<<1,1>>>:

• Single thread enforces the sequential inter-batch threshold constraint (running sum of seq_lens_this_time across batch items)
• Copies matched draft tokens from NgramMatchResult scratch buffer to output tensors
• This serial phase is necessary because batch k's draft token budget depends on batches 0..k-1's finalized results

Shared device code (ngram_match_common.cuh):

• NgramMatchResult struct — inter-phase communication via device memory scratch buffer
• atomicMin64() — 64-bit CAS device function for leftmost-match atomics
• parallel_ngram_search() — block-cooperative sliding-window search used by both kernels

File Changes

New shared header (1 file):

• ngram_match_common.cuh: NgramMatchResult, atomicMin64(), parallel_ngram_search() device functions. No __global__ kernels in the header (avoids multiple-definition linker errors).

CUDA kernels (2 files):

• ngram_match.cu: Two __global__ kernels (ngram_match_search_kernel + ngram_match_gather_kernel). Host function NgramMatch() launches Phase 1 <<<max_batch_size, 256, 0, stream>>> then Phase 2 <<<1, 1, 0, stream>>>. Uses seq_lens_encoder / seq_lens_decoder.
• ngram_match_mixed.cu: Two __global__ kernels (ngram_match_mixed_search_kernel + ngram_match_mixed_gather_kernel). Host function HybridMtpNgram() launches Phase 1 then Phase 2. Uses seq_lens_this_time / seq_lens_decoder. Gather kernel computes ori_seq_len_this_time per-batch.

Python callers (2 files):

• ngram.py: Removed ~10 .cpu() tensor copies in _run_impl(). All tensors stay on device.
• mtp.py: Removed .cpu()/.cuda() round-trips and CUDAPinnedPlace copy in _extend_draft_token_with_ngram_match().

Design Decisions

1. Why two-phase (not fully parallel)?

The CPU kernels maintain a running threshold sum across batch items: each batch's seq_lens_this_time[i] affects the draft token budget for subsequent batches. This is a data-dependent sequential dependency — batch k cannot finalize until batches 0..k-1 have computed their match results.

Approach	Description	Verdict
Two-phase (search ∥ gather serial)	Phase 1: all batches search in parallel. Phase 2: single thread applies threshold + copies tokens	Chosen — parallelizes the expensive O(bsz × seq_len) search while preserving exact semantics
Fully serial <<<1,1>>>	1 thread processes all batches sequentially	Rejected — reviewer feedback: not utilizing GPU parallelism for bsz=256, seq_len=128k
Prefix-sum + parallel search	Compute threshold via parallel scan, then parallel gather	Rejected — threshold depends on match RESULTS (data-dependent), not just input

2. atomicMin64 for leftmost-match

Multiple threads in a block may find valid ngram matches at different positions. The leftmost match must win (matching CPU semantics). We use a 64-bit Compare-And-Swap loop (atomicCAS on unsigned long long) to atomically update the minimum match position without locks.

3. Kernel differences: ngram_match vs ngram_match_mixed

Both kernels call the same parallel_ngram_search() device function. Business-specific differences:

Aspect	ngram_match	ngram_match_mixed
write_offset	1	ori_seq_len_this_time
min_ngram_size	1 (fixed)	Configurable
Default threshold	128 (INFER_WITH_REFERENCE_TOKENUM_THRESHOLD)	1024 (SPEC_TOKENUM_THRESHOLD)
Batch-skip condition	seq_lens_encoder > 0	ori_seq_len_this_time == 0

4. Zero-copy memory access

Before (CPU path): 10 D2H + 3 H2D copies per call, each triggering cudaStreamSynchronize.
After (CUDA path): All tensors stay on device. Net: 13 sync points → 0.

Usage or Command

No API changes. The CUDA kernels are drop-in replacements — same function signatures, same op registration, same Python call sites.

# Build FastDeploy (ops are compiled automatically)
bash build.sh

# Run correctness + latency tests
python -m pytest tests/spec_decode/test_ngram_gpu_kernel.py -v

# Existing speculative decoding workflows work unchanged:
python -m fastdeploy.entrypoints.openai.api_server \
    --model baidu/ERNIE-4.5-21B-A3B-Paddle \
    --speculative_method ngram

Accuracy Tests

CI environment: SM90 H20 GPU, CUDA 12.6, Python 3.10 (run_tests_with_coverage job).

All 11 tests passed (+ 8 subtests) in 101.44s:

Correctness Tests (NgramMatch kernel)

Test	Config	Result
test_correctness_basic	bsz=4, seeds vary	PASSED
test_correctness_varied_seeds	seeds=0,7,123,999	4/4 PASSED
test_large_batch_long_seq	bsz=256, input_len=131072	PASSED
test_many_short_seqs	bsz=256, input_len=1024	PASSED
test_single_batch_long_seq	bsz=1, seq_len=128k	PASSED

Correctness Tests (HybridMtpNgram kernel)

Test	Config	Result
test_correctness_basic	bsz=4, seeds vary	PASSED
test_correctness_varied_seeds	seeds=0,7,123,999	4/4 PASSED
test_large_batch_long_seq	bsz=256, input_len=131072	PASSED
test_many_short_seqs	bsz=256, input_len=1024	PASSED
test_single_batch_long_seq	bsz=1, seq_len=128k	PASSED

Latency Benchmark (CI-verified, SM90 H20)

Metric	GPU kernel (zero-copy)	CPU path (with D2H/H2D)
Per-call latency (batch=32, input_len=512, 100 runs)	0.690 ms	0.953 ms
Speedup	1.38×	baseline
CUDA sync points per call	0	13

Existing operator tests also passed:

• test_ngram_match.py::TestNgramMatchOp::test_basic_match ✅
• test_ngram_match.py::TestNgramMatchOp::test_no_match ✅
• test_hybrid_mtp_ngram.py::TestNgramMatchMixed::test_ngram_match_mixed ✅

Checklist

• Two-phase parallel CUDA kernel (<<<bsz, 256>>> search + <<<1,1>>> gather)
• atomicMin64 CAS for leftmost-match semantics
• Tested at reviewer-specified scale: bsz=256, seq_len=128k
• CI-verified: 11/11 tests passed on SM90 H20 (101.44s)
• Latency benchmark: 1.38× speedup (GPU 0.690ms vs CPU 0.953ms)
• Existing operator tests pass (test_ngram_match, test_hybrid_mtp_ngram)
• No API changes (drop-in replacement)
• pre-commit hooks pass (black, isort, clang-format, flake8, ruff)

[paddle-bot]

Thanks for your contribution!

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (develop@0b4c1cb). Learn more about missing BASE report.

Additional details and impacted files

@@            Coverage Diff             @@
##             develop    #6960   +/-   ##
==========================================
  Coverage           ?   73.08%           
==========================================
  Files              ?      402           
  Lines              ?    56419           
  Branches           ?     8903           
==========================================
  Hits               ?    41236           
  Misses             ?    12272           
  Partials           ?     2911           

Flag	Coverage Δ
GPU	73.08% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cloudforge1]

@luotao1 CI green — 35/35 checks passed (HPU/iluvatar infra-only failures). 5/5 kernel tests passed on SM90 H20, GPU 0.934ms vs CPU 0.965ms (1.03×, 13→0 sync points). @CSWYF3634076 ready for review.

[cloudforge1]

@luotao1 PR #6960（No.49 ngram_match CUDA优化）CI 已通过 — SM90 实测 1.03× 加速，同步点 13→0。请问由哪位评审？

[freeliuzc]

改为 cuda Kernel 不是简单的把逻辑改为 cuda，而是需要用并行策略加速 Kernel 哈，比如最大会有 bsz=256，seq_len=128k

[cloudforge1]

感谢指出，当前 <<<1,1>>> 确实没有利用 GPU 并行能力。

已着手在本 PR 内重构为并行版本，初步方案：

• batch 维度：grid=bsz，每个 batch item 一个 block
• seq_len 维度：block 内多线程并行滑窗搜索 ngram 匹配位置（适配 bsz=256, seq_len=128k）
• 跨 batch 的 running sum 依赖通过分阶段处理或 prefix scan 解耦

这个方向是否符合预期？更新后会补充大 batch 场景的性能对比数据。

[cloudforge1]

已完成并行重构，CI 已通过（SM90 H20）。

架构：两阶段 kernel

• Phase 1 <<<bsz, 256>>>：每个 block 处理一个 batch item，256 线程并行滑窗搜索 + atomicMin64 CAS 保证最左匹配
• Phase 2 <<<1,1>>>：串行 threshold 约束（跨 batch 依赖）+ token 拷贝

CI 测试结果（11/11 passed，101.44s）：

• test_large_batch_long_seq：bsz=256, input_len=131072 — 两个 kernel 均通过
• latency benchmark（batch=32, input_len=512, 100 runs）：GPU 0.690ms vs CPU 0.953ms = 1.38× 加速
• 13 个 D2H/H2D 同步点 → 0

共享设备代码在 ngram_match_common.cuh（NgramMatchResult struct + parallel_ngram_search()），两个 kernel 复用相同搜索逻辑。