add fused mhc_post_pre kernel

[gnovack]

Purpose

This PR adds a new mHC kernel, which fuses the hc_post operation with the prenorm_gemm portion of hc_pre. The approach is adapted from TRTLLM, and performs the GEMM using FMA (rather than tensor cores), improving speed at low concurrency.

Benchmarks

Benchmark results with deepseek-ai/DeepSeek-V4-Flash at concurrency 4.

Before this PR:

================= Serving Benchmark Result =================
Successful requests:                     300       
Failed requests:                         0         
Maximum request concurrency:             4         
Benchmark duration (s):                  302.63    
Total input tokens:                      307200    
Total generated tokens:                  100753    
Request throughput (req/s):              0.99      
Output token throughput (tok/s):         332.92    
Peak output token throughput (tok/s):    408.00    
Peak concurrent requests:                7.00      
Total token throughput (tok/s):          1348.00   
--------------------Time to First Token---------------------
Mean TTFT (ms):                          196.89    
Median TTFT (ms):                        192.03    
P99 TTFT (ms):                           680.64    
P90 TTFT (ms):                           206.13    
----------Time per Output Token (excl. 1st token)-----------
Mean TPOT (ms):                          11.40     
Median TPOT (ms):                        11.34     
P99 TPOT (ms):                           14.48     
P90 TPOT (ms):                           12.28
============================================================

With this PR:

================= Serving Benchmark Result =================
Successful requests:                     300       
Failed requests:                         0         
Maximum request concurrency:             4         
Benchmark duration (s):                  284.77    
Total input tokens:                      307200    
Total generated tokens:                  100436    
Request throughput (req/s):              1.05      
Output token throughput (tok/s):         352.69    
Peak output token throughput (tok/s):    432.00    
Peak concurrent requests:                7.00      
Total token throughput (tok/s):          1431.47   
--------------------Time to First Token---------------------
Mean TTFT (ms):                          185.48    
Median TTFT (ms):                        188.37    
P99 TTFT (ms):                           539.61    
P90 TTFT (ms):                           201.03    
----------Time per Output Token (excl. 1st token)-----------
Mean TPOT (ms):                          10.74     
Median TPOT (ms):                        10.69     
P99 TPOT (ms):                           12.39     
P90 TPOT (ms):                           11.59
============================================================

Accuracy Evals

task	result
gsm8k	flexible-extract=0.9492; strict-match=0.9500
aime25	exact_match=0.9333

Still Pending

• Code cleanup
• Rebase
• Add unit tests
• Add model quality eval results to PR
• Look into autotuning to avoid hard-coded metaparams (tile_n, split_k, etc.)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gnovack.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces Multi-Head-Channel (MHC) fused kernels, including a fused pmap and GEMM operation, to optimize the DeepSeek-V4 model. The feedback highlights a missing source file for the mhc_fused_hc binding, an inconsistency in token count thresholds for kernel selection, and the use of non-flattened tensors in a kernel call. Additionally, a potential UnboundLocalError was identified in the model's forward loop when no layers are processed.

[gemini-code-assist]

The else block of the for loop uses the layer variable, which will raise an UnboundLocalError if the loop is empty (e.g., if start_layer == end_layer). This can happen in certain tensor parallel or pipeline parallel configurations. A check should be added to ensure layer is defined before use, or better, check if residual is not None to ensure at least one layer was processed.

Suggested change

	else:
	hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
	if residual is not None:
	hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[zyongye]

If tilelang kernel is here to stay, can we add warmup logic in the engine? So that it won't trigger compile in realtime. It should be similar to deepgemm and flashinfer