Implementation:Sgl project Sglang Spatial Partitioning

Knowledge Sources	Sgl_project_Sglang
Domains	GPU Computing, CUDA Streams, Spatial Multiplexing
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for CUDA spatial SM (Streaming Multiprocessor) partitioning using green contexts, enabling concurrent kernel execution on partitioned GPU resources.

Description

This module provides the spatial partitioning API for the SGLang kernel library. It wraps low-level CUDA green context operations exposed through the spatial_ops C++ extension registered via PyTorch's custom operator mechanism.

The primary function, create_greenctx_stream_by_value, creates two CUDA ExternalStream objects backed by green contexts with specified SM allocations. Each stream is assigned a distinct partition of the GPU's Streaming Multiprocessors, allowing kernels dispatched on different streams to execute concurrently on non-overlapping hardware resources. The companion function, get_sm_available, queries the total SM count on a given device so callers can compute appropriate partition sizes.

The module lazily imports the spatial_ops extension at module load time and raises an ImportError with a descriptive message if the extension is unavailable (typically requiring CUDA Driver >= 12.4).

Usage

Use this module when you need to spatially multiplex GPU SMs between different workloads during inference, such as running attention computation and MoE (Mixture of Experts) computation concurrently on separate SM partitions. First call get_sm_available to determine the total SM count, then call create_greenctx_stream_by_value to create two partitioned streams with the desired SM allocation.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/spatial.py
• Lines: 1-64

Signature

def create_greenctx_stream_by_value(
    SM_a: int, SM_b: int, device_id: int = None
) -> tuple[ExternalStream, ExternalStream]:
    ...

def get_sm_available(device_id: int = None) -> int:
    ...

Import

from sgl_kernel import create_greenctx_stream_by_value, get_sm_available

I/O Contract

Inputs

create_greenctx_stream_by_value

Name	Type	Required	Description
SM_a	int	Yes	Number of Streaming Multiprocessors to allocate to stream A
SM_b	int	Yes	Number of Streaming Multiprocessors to allocate to stream B
device_id	int	No	CUDA device ID; defaults to the current device via torch.cuda.current_device()

get_sm_available

Name	Type	Required	Description
device_id	int	No	CUDA device ID; defaults to the current device via torch.cuda.current_device()

Outputs

create_greenctx_stream_by_value

Name	Type	Description
(stream_a, stream_b)	tuple[ExternalStream, ExternalStream]	Two CUDA ExternalStream objects backed by green contexts, each allocated the specified number of SMs

get_sm_available

Name	Type	Description
sm_count	int	Total number of Streaming Multiprocessors available on the specified device

Usage Examples

import torch
from sgl_kernel import create_greenctx_stream_by_value, get_sm_available

# Query the total number of SMs on the current device
total_sms = get_sm_available()

# Split SMs: 80% for attention, 20% for MoE
sm_attention = int(total_sms * 0.8)
sm_moe = total_sms - sm_attention

# Create two partitioned CUDA streams
stream_attention, stream_moe = create_greenctx_stream_by_value(sm_attention, sm_moe)

# Run attention kernel on stream_attention
with torch.cuda.stream(stream_attention):
    # ... launch attention kernels here ...
    pass

# Run MoE kernel concurrently on stream_moe
with torch.cuda.stream(stream_moe):
    # ... launch MoE kernels here ...
    pass

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime