minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix Inkling kernel imports after migration (#32076)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-07-22 22:09:01 +08:00
committed by GitHub
parent 0a6d1930c3
commit b855efd9e6
2 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/inkling_common/attn.py
+2 -1
View File
@@ -364,7 +364,8 @@ class InklingAttention(nn.Module):
def _fused_attn_prologue_verify(self, q, k, v, forward_batch, log_scaling_tau=None): def _fused_attn_prologue_verify(self, q, k, v, forward_batch, log_scaling_tau=None):
"""Fused target-verify {k/v sconv + save_windows + qk-norm (+ KV store)} """Fused target-verify {k/v sconv + save_windows + qk-norm (+ KV store)}
(jit_kernel/inkling_attn_prologue.py); returns ``(q, k, v, did_store)``. (kernels/ops/model/inkling/inkling_attn_prologue.py); returns
``(q, k, v, did_store)``.
The fused kernel writes raw bf16 KV, so it only does the store when the The fused kernel writes raw bf16 KV, so it only does the store when the
KV pool is bf16: full layers at ``out_cache_loc`` in the full pool, KV pool is bf16: full layers at ``out_cache_loc`` in the full pool,
python/sglang/srt/models/inkling_common/kernels/comm.py
+7 -5
View File
@@ -94,7 +94,7 @@ def _ar_jit():
lazy so importing comm.py doesn't pull in the JIT machinery).""" lazy so importing comm.py doesn't pull in the JIT machinery)."""
if not is_cuda(): if not is_cuda():
return None return None
from sglang.jit_kernel import inkling_all_reduce from sglang.kernels.ops.model.inkling import inkling_all_reduce
return inkling_all_reduce return inkling_all_reduce
@@ -103,7 +103,7 @@ def _ar_jit():
def _ar_fused_jit(): def _ar_fused_jit():
if not is_cuda(): if not is_cuda():
return None return None
from sglang.jit_kernel import inkling_ar_fused from sglang.kernels.ops.model.inkling import inkling_ar_fused
return inkling_ar_fused return inkling_ar_fused
@@ -240,7 +240,8 @@ def ar_sconv_norm_fusable(
"""True when a decode {all-reduce -> sconv -> add+RMSNorm} chain """True when a decode {all-reduce -> sconv -> add+RMSNorm} chain
(attn-side: wo_ud AR -> attn_sconv -> mlp_norm; MoE-side: MoE AR -> (attn-side: wo_ud AR -> attn_sconv -> mlp_norm; MoE-side: MoE AR ->
mlp_sconv -> next attn_norm) mlp_sconv -> next attn_norm)
can run as the single fused kernel (jit_kernel/inkling_ar_fused.py). Must be can run as the single fused kernel
(kernels/ops/model/inkling/inkling_ar_fused.py). Must be
evaluated identically by the producing layer (MoE ``reduce=False``) and the evaluated identically by the producing layer (MoE ``reduce=False``) and the
consuming layer/tail -- it is a pure function of per-forward state.""" consuming layer/tail -- it is a pure function of per-forward state."""
if not is_cuda(): if not is_cuda():
@@ -674,7 +675,7 @@ def all_gather_hidden(input: torch.Tensor, group: GroupCoordinator) -> torch.Ten
def _ar_ssconv_jit(): def _ar_ssconv_jit():
if not is_cuda(): if not is_cuda():
return None return None
from sglang.jit_kernel import inkling_ar_scattered_sconv from sglang.kernels.ops.model.inkling import inkling_ar_scattered_sconv
return inkling_ar_scattered_sconv return inkling_ar_scattered_sconv
@@ -688,7 +689,8 @@ def scattered_ar_sconv_fusable(
) -> bool: ) -> bool:
"""True when an extend {reduce_scatter_hidden -> sconv(shard) -> """True when an extend {reduce_scatter_hidden -> sconv(shard) ->
all_gather_hidden} chain can run as the single fused v3/v3b-style kernel all_gather_hidden} chain can run as the single fused v3/v3b-style kernel
(jit_kernel/inkling_ar_scattered_sconv.py). Pure function of per-forward (kernels/ops/model/inkling/inkling_ar_scattered_sconv.py). Pure function of
per-forward
state -- the producing layer (reduce=False) and the consuming site must state -- the producing layer (reduce=False) and the consuming site must
evaluate it identically.""" evaluate it identically."""
if not is_cuda(): if not is_cuda():