[Kernels] add triton contiguous groupgemm

[lessw2020]

This PR adds a triton based contiguous group gemm and integrates it with deepseek.
This group gemm has full backwards support.

Testing:
a - have verified cg group gemm works for deepseek inference.
b - have added and run unit tests for both forward and backward. (The perf test is against manual looping).

===== Testing Backward Pass: DeepSeek Shapes   =====
Testing shape: M=512, K=128, N=128, group_size=128
Testing shape: M=1,024, K=1,024, N=1,024, group_size=128
.
===== Testing Forward Pass: DeepSeek Shapes =====
Testing shape: M=4,096, K=1,024, N=1,024, group_size=128
Testing shape: M=4,096, K=4,096, N=4,096, group_size=128
Testing shape: M=2,048, K=4,096, N=7,168, group_size=128
Testing shape: M=2,048, K=7,168, N=2,048, group_size=128
.
===== Benchmarking Backward Pass with DeepSeek Shapes   =====
Benchmarking shape: M=512, K=128, N=128, group_size=128
Benchmarking shape: M=1,024, K=1,024, N=1,024, group_size=128

DeepSeek Backward Performance Summary:
Shape                          CG-GEMM (ms)    PyTorch (ms)    Speedup    TFLOPS    
--------------------------------------------------------------------------------
M=512, K=128, N=128            0.305           0.837           2.74       0.11      
M=1024, K=1024, N=1024         0.304           1.293           4.25       14.12     
.
===== Benchmarking Forward Pass with DeepSeek Shapes =====
Benchmarking shape: M=2,048, K=1,024, N=1,024, group_size=128
Benchmarking shape: M=4,096, K=4,096, N=4,096, group_size=128

DeepSeek Forward Performance Summary:
Shape                          CG-GEMM (ms)    PyTorch (ms)    Speedup    TFLOPS    
--------------------------------------------------------------------------------
M=2048, K=1024, N=1024         0.050           0.617           12.46      86.69     
M=4096, K=4096, N=4096         0.380           1.569           4.13       361.60    
.
----------------------------------------------------------------------
Ran 4 tests in 6.268s

OK

[lessw2020]

GPU test failure is not related to this PR:
rank0]:[rank0]: RuntimeError: CUDA driver error: invalid device context

[kwen2501]

The integration code looks good to me.
Didn't look at the kernel closely. Did you check if the generated results are correct?
If correct, approving to unblock.

[lessw2020]

The integration code looks good to me. Didn't look at the kernel closely. Did you check if the generated results are correct? If correct, approving to unblock.

thanks and yes, verified results via unit tests (variety of shapes) as well as generate/inference producing valid output.

[yuan-luo]

Triton v3.4.0 doesn't support fill_1d_tma_descriptor and fill_2d_tma_descriptor, only support fill_tma_descriptor. Do you have any plan to refactor the code to support TMA in new Triton?

/opt/conda/lib/python3.10/site-packages/sglang/srt/layers/moe/fused_moe_ibm/persistent_gg_bf16_tma.py:291: in grouped_gemm_persistent
    return _grouped_gemm_persistent(x, w, expert_indices)
/opt/conda/lib/python3.10/site-packages/sglang/srt/layers/moe/fused_moe_ibm/persistent_gg_bf16_tma.py:196: in _grouped_gemm_persistent
    desc_helper = TmaDescriptorHelper(tma_size=tma_size)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <sglang.srt.layers.moe.fused_moe_ibm.tma_autotune.TmaDescriptorHelper object at 0x7f1975b901c0>, tma_size = 128

    def __init__(self, tma_size: int = 128):
        """Initialize the TMA descriptor helper.
    
        Args:
            tma_size: Size of the TMA descriptor in bytes
        """
        if not CudaUtils.verify_tma():
            raise RuntimeError(
                "TMA not supported on this device (requires Hopper or newer)"
            )
    
        self.tma_size = tma_size
>       self.fill_1d_tma_descriptor_inner = driver.active.utils.fill_1d_tma_descriptor
E       AttributeError: 'CudaUtils' object has no attribute 'fill_1d_tma_descriptor'. Did you mean: 'fill_tma_descriptor'?

[lessw2020]

Hi @yuan-luo - I'm not working in Triton anymore (hence didn't realize they changed the descriptors) but I think @AdnanHoque can quickly make the update.