[jit-kernel] Add CuTe DSL GDN Decode Kernel

[liz-badada]

Co-author: @zhou9402, @HongliMi, @xutizhou

Motivation

Modifications

Add CuTe DSL GDN Decode Kernel

• hybrid_linear_attn_backend.py - Add CuTe DSL decode branch (env: SGLANG_USE_CUTEDSL_GDN_DECODE=1)
• python/sglang/jit_kernel/cutedsl_gdn.py - CuTe DSL kernel to support decode gdn
• python/sglang/jit_kernel/tests/test_cutedsl_gdn.py - Precision & performance tests

nvidia-cutlass-dsl >=4.3.0

Note: This is a ssm state transpose free version, will also have a version with better performance (need pre-transpose ssm state just once) @HongliMi.

Accuracy Tests

pytest path/to/python/sglang/jit_kernel/tests/test_cutedsl_gdn.py -v -s

SGLANG_USE_CUTEDSL_GDN_DECODE=1 python -m sglang.bench_one_batch --model-path Qwen/Qwen3-Next-80B-A3B-Instruct \
    --tp 2 --batch 1 --input-len 128 --output-len 64 --cuda-graph-bs 1 --correctness-test

Benchmarking and Profiling

SGLANG_USE_CUTEDSL_GDN_DECODE=1 python3 -m sglang.bench_one_batch_server --model-path Qwen/Qwen3-Next-80B-A3B-Instruct \
    --batch-size 128 --input-len 128 1024 2048 4096 8192 --output-len 1024 --disable-radix-cache \
    --tp 2 --cuda-graph-bs 128 --show-report

E2E (H200, TP=2, Output lens: [1024], E2E output throughput speedup: 4.6% - 5.2%)

Triton

batch size	input len	latency (s)	input throughput (tok/s)	output throughput (tok/s)	acc length	ITL (ms)	input cost ($/1M)	output cost ($/1M)	cache hit rate
128	1024	25.87	56064.6	5568.82	n/a	22.99	0.03	0.2	n/a
128	2048	26.74	56155.4	5938.23	n/a	21.56	0.03	0.19	n/a
128	4096	31.56	55581.9	5924.45	n/a	21.61	0.03	0.19	n/a
128	8192	42.45	53994	5691.37	n/a	22.49	0.03	0.2	n/a

CuTe DSL

batch size	input len	latency (s)	input throughput (tok/s)	output throughput (tok/s)	acc length	ITL (ms)	input cost ($/1M)	output cost ($/1M)	cache hit rate
128	1024	24.85	55840.9	5825.61	n/a	21.97	0.03	0.19	n/a
128	2048	25.73	56161.4	6223.17	n/a	20.57	0.03	0.18	n/a
128	4096	30.54	55649.4	6206.95	n/a	20.62	0.03	0.18	n/a
128	8192	41.26	54103.8	5989.78	n/a	21.37	0.03	0.19	n/a

E2E (B200, TP=2, Output lens: [1024], E2E output throughput speedup: 2.6% - 3.4%)

Triton

batch size	input len	latency (s)	input throughput (tok/s)	output throughput (tok/s)	acc length	ITL (ms)	input cost ($/1M)	output cost ($/1M)	cache hit rate
128	1024	20.11	84113.8	7064.7	n/a	18.12	0.04	0.31	n/a
128	2048	20.99	83395.6	7344.68	n/a	17.43	0.04	0.3	n/a
128	4096	24.57	80808.5	7249.77	n/a	17.66	0.04	0.31	n/a
128	8192	33.29	76240.3	6709.32	n/a	19.08	0.04	0.33	n/a

CuTe DSL

batch size	input len	latency (s)	input throughput (tok/s)	output throughput (tok/s)	acc length	ITL (ms)	input cost ($/1M)	output cost ($/1M)	cache hit rate
128	1024	19.51	83999.9	7301.8	n/a	17.53	0.04	0.3	n/a
128	2048	20.48	83453.1	7558.23	n/a	16.94	0.04	0.29	n/a
128	4096	24.1	81450.2	7420.56	n/a	17.25	0.04	0.3	n/a
128	8192	32.79	76249.7	6884.49	n/a	18.59	0.04	0.32	n/a

E2E (H20, TP=2, Output lens: [1024], E2E output throughput speedup: 1.7% - 2.5%)

Triton

batch size	input len	latency (s)	input throughput (tok/s)	output throughput (tok/s)	acc length	ITL (ms)	input cost ($/1M)	output cost ($/1M)	cache hit rate
128	1024	35.74	22943.5	4365.37	n/a	29.32	0.07	0.25	n/a
128	2048	40.33	22939.3	4534.73	n/a	28.23	0.07	0.25	n/a
128	4096	52.35	22592.1	4497.54	n/a	28.46	0.07	0.25	n/a
128	8192	79.48	21679.9	4212.54	n/a	30.39	0.07	0.26	n/a

CuTe DSL

batch size	input len	latency (s)	input throughput (tok/s)	output throughput (tok/s)	acc length	ITL (ms)	input cost ($/1M)	output cost ($/1M)	cache hit rate
128	1024	35.24	22992.2	4437.69	n/a	28.84	0.07	0.25	n/a
128	2048	39.62	22955.2	4647.63	n/a	27.54	0.07	0.24	n/a
128	4096	51.86	22568	4578.15	n/a	27.96	0.07	0.24	n/a
128	8192	78.83	21665.1	4307.04	n/a	29.72	0.07	0.26	n/a

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.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

Hello @liz-badada, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new, highly optimized CuTe DSL kernel for the Fused Sigmoid Gating Delta Rule (GDN) decode operation, providing an alternative to the existing Triton-based implementation. The primary goal is to enhance performance for GDN decode by leveraging advanced GPU programming features like CP.ASYNC and optimized memory access patterns. The new kernel is integrated with a toggle for easy activation and is thoroughly tested for both accuracy and speed against the current reference.

Highlights

• New CuTe DSL GDN Decode Kernel: Introduced a high-performance CuTe DSL implementation for the Fused Sigmoid Gating Delta Rule (GDN) decode operation, designed for efficiency on NVIDIA GPUs.
• Conditional Kernel Selection: The new CuTe DSL kernel is integrated into the hybrid_linear_attn_backend.py and can be enabled via the SGLANG_USE_CUTEDSL_GDN_DECODE=1 environment variable, allowing users to switch between the new CuTe DSL kernel and the existing Triton kernel.
• Optimized Kernel Design: The CuTe DSL kernel leverages advanced GPU programming techniques such as CP.ASYNC for efficient global-to-shared memory transfers, bank-conflict-free access patterns with swizzle layouts, and automatic batch size selection (small vs. big batch kernels) for optimal performance.
• Comprehensive Testing: New tests have been added to verify both the precision and performance of the CuTe DSL GDN kernel against the Triton reference implementation, covering different batch sizes (B=16 and B=128).
• Dependency Update: The nvidia-cutlass-dsl dependency has been updated to version >=4.3.0 to support the new CuTe DSL kernel features.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new, high-performance CuTe DSL kernel for the GDN decode operation, which can be enabled via an environment variable. The changes are well-structured, including the kernel implementation, its integration into the existing attention backend with lazy loading, and a comprehensive test suite for precision and performance against the existing Triton kernel. My feedback focuses on improving code maintainability by reducing code duplication in the kernel selection logic and enhancing clarity and consistency in the test suite. Overall, this is a solid contribution.

[Fridge003]

@hebiao064 Can you please take a look

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/tag-and-rerun-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[BBuf]

CI failed: https://github.com/sgl-project/sglang/actions/runs/21064410750/job/60608238248?pr=15631

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[liz-badada]

/rerun-failed-ci

[BBuf]

@liz-badada Can you have a look at this error: https://github.com/sgl-project/sglang/actions/runs/22986788908/job/66756957189?pr=20395