WIP: kdn decode pto-isa

[fiskrt]

A vector-only PTO-ISA kernel for the kimi-delta-attention (KDA) decoding (extend in sglang lingo) stage. Will mostly be called from sglang in packed format (i.e B=1, T=number of sequences and cu_seqlens are just 1's), initial state is always provided along with a list of the state indices. For kimi-linear l2norm=True.

TODO:

• fuse the gating into kernel
• benchmark script (compare to roofline 1.5 TB/s 910B2 and current Triton kernel)
• A5 supprt
• Kimi-K3 support

Bench

Kernel vs kernel

Benchmark against the Triton kernel. N/a means that Triton fails due to too many blocks launched (max is 64k)

• We see a ~50% gain across the board

End to end

Depending on the context length the kdn doce kernel is can between 5-20% of a linear layer, and around 1-8% of whole transformer.

At 128 concurrent decoding sessions , pto-isa sees avg 125 us, while triton sees 160 us. With short sequence length at this setting this gives us an end2end speedup of 5%. Measuring with sglang.bench_one_batch to avoid the scheduler overhead.

KDA_DECODE_PTO_BACKEND=0 ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 SGLANG_TORCH_PROFILER_DIR=./prof_output_tri/ python -m sglang.bench_one_batch   --model-path kimi-linear/ --trust-remote-code   --tp 4 --context-length 1024 --mem-fraction-static 0.95   --batch-size 128 --input-len 1 --output-len 4 --profile

PTO:

Triton

Notes:

• Sglang has support for calling recurrent implementation for pre-fill aswell in some range? Maybe up to 31/63 tokens? Ours work for unpacked T>1, and packed cu_seqlens[i+1]-cu_seqlens[i]>1 aswell but the double buffering and testing could be more looked in to.
• Current triton kernel is wrong assumes tranposed state layout?

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.