Kernel: optimize decoding metadata in NSA multi-spec backend with fused kernels

[Johnsonms]

Motivation

Implement fused CUDA kernels to eliminate redundant metadata copies in Native Sparse Attention (NSA) backend during CUDA graph replay for speculative decoding. This optimization provides 3-5x speedup for multi-backend metadata operations.

Changes

Core Implementation

• Add fused_metadata_copy_cuda: Single-backend fused kernel supporting DECODE, TARGET_VERIFY, and DRAFT_EXTEND forward modes
• Add fused_metadata_copy_multi_cuda: Multi-backend kernel that copies metadata to 3 backends simultaneously in a single kernel launch

Runtime Optimization

• Update nsa_backend.py to intelligently use fused kernels:
• speculative_num_steps >= 3: Use fused kernel for first 3 backends,
  then copy remaining backends individually
• speculative_num_steps < 3: Use individual copy (overhead not worth it)

Testing

• Create comprehensive test suite covering:
• Single-backend kernel: all forward modes, optional tensors
• Multi-backend kernel: 3-backend simultaneous copy
• Performance benchmarks with timing and speedup measurements

Performance Impact

CPU side

Before:

After:

GPU side

Before

After

Performance Improvements

CPU-side metadata processing:

• Before: 193μs (24 kernel launches)
• After: 12μs (1 fused kernel launch)
• Speedup: 16x faster (193μs → 12μs)

Kernel execution time:

• Before: 177μs
• After: 4.7μs
• Speedup: 37.6x faster (177μs → 4.7μs)

End-to-end throughput:

• Before: 151 tokens/sec
• After: 179 tokens/sec
• Improvement: +18.5% (+28 tokens/sec)

Technical Details

The fused kernels handle:

• Basic metadata: cache_seqlens, cu_seqlens_k, page_table, nsa_cache_seqlens
• Optional tensors: real_page_table, flashmla_num_splits, flashmla_metadata
• All copies done in single kernel launch, reducing PCIe/memory overhead

Files Modified

• sgl-kernel/csrc/elementwise/fused_metadata_copy.cu: CUDA kernel impl
• sgl-kernel/python/sgl_kernel/elementwise.py: Python bindings
• python/sglang/srt/layers/attention/nsa_backend.py: Runtime integration
• sgl-kernel/tests/test_fused_metadata_copy.py: Comprehensive tests

Tested: All 40 tests passing (100% pass rate)

Modifications

Accuracy Tests

1. Accuracy Test with gsm8k

python3 -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3.2-Exp \
  --trust-remote-code \
  --tp-size 8 --dp-size 8 --enable-dp-attention \
  --tool-call-parser deepseekv31 \
  --reasoning-parser deepseek-v3 \
  --chat-template ./examples/chat_template/tool_chat_template_deepseekv32.jinja

python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319

2. Accuracy Test with gpqa-diamond
   Service:
   python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4

python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa --num-examples 198 --max-tokens 120000 --repeat 8 --thinking-mode deepseek-v3

3. Accuracy Test with aime 2025
   python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --tool-call-parser deepseekv32 --reasoning-parser deepseek-v3

#! /bin/bash
export NEMO_SKILLS_DISABLE_UNCOMMITTED_CHANGES_CHECK=1

ns prepare_data aime25

PORT=30000
BACKEND=sglang
MODEL="deepseek-ai/DeepSeek-V3.2-Exp" # Should be changed to the model name
MODEL_NAME="dsv32-fp8"

echo "Starting AIME25 evaluation with model $MODEL on port $PORT using backend $BACKEND..."
ns eval \
  --benchmarks=aime25:4 \
  --server_type=$BACKEND \
  --model=$MODEL \
  --server_address=http://localhost:${PORT}/v1 \
  --output_dir=nemo_skills_aime25_${MODEL_NAME}_output_${BACKEND}_$(date +%Y%m%d_%H%M%S) \
  ++chat_template_kwargs.thinking=true \
  ++inference.temperature=1.0 \
  ++inference.top_p=0.95 \
  ++inference.tokens_to_generate=64000
  # ++inference.tokens_to_generate=120000 for Speciale model

4. Correctness
   After the above tests, using both kernels and verified every steps

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Fridge003]

Can you please move the kernels to jit folder, thanks~

[Johnsonms]

Re-performed the accuracy testing

Accuracy Tests

1. Accuracy Test with gsm8k

python3 -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3.2-Exp \
  --trust-remote-code \
  --tp-size 8 --dp-size 8 --enable-dp-attention \
  --tool-call-parser deepseekv31 \
  --reasoning-parser deepseek-v3 \
  --chat-template ./examples/chat_template/tool_chat_template_deepseekv32.jinja

python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319

2. Accuracy Test with gpqa-diamond
   Service:
   python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4
python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa --num-examples 198 --max-tokens 120000 --repeat 8 --thinking-mode deepseek-v3

3. Accuracy Test with aime 2025
   python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --tool-call-parser deepseekv32 --reasoning-parser deepseek-v3

#! /bin/bash
export NEMO_SKILLS_DISABLE_UNCOMMITTED_CHANGES_CHECK=1

ns prepare_data aime25

PORT=30000
BACKEND=sglang
MODEL="deepseek-ai/DeepSeek-V3.2-Exp" # Should be changed to the model name
MODEL_NAME="dsv32-fp8"

echo "Starting AIME25 evaluation with model $MODEL on port $PORT using backend $BACKEND..."
ns eval \
  --benchmarks=aime25:4 \
  --server_type=$BACKEND \
  --model=$MODEL \
  --server_address=http://localhost:${PORT}/v1 \
  --output_dir=nemo_skills_aime25_${MODEL_NAME}_output_${BACKEND}_$(date +%Y%m%d_%H%M%S) \
  ++chat_template_kwargs.thinking=true \
  ++inference.temperature=1.0 \
  ++inference.top_p=0.95 \
  ++inference.tokens_to_generate=64000
  # ++inference.tokens_to_generate=120000 for Speciale model

[Johnsonms]

Can you please move the kernels to jit folder, thanks~

Done, with re-performance Accuracy Tests. Thanks @Fridge003 !

[DarkSharpness]

/tag-and-rerun-ci

[yuan-luo]

Awesome job!