diff --git a/python/sglang/srt/models/inkling_common/attn.py b/python/sglang/srt/models/inkling_common/attn.py index 5a2f23b48..cd2ab443c 100644 --- a/python/sglang/srt/models/inkling_common/attn.py +++ b/python/sglang/srt/models/inkling_common/attn.py @@ -364,7 +364,8 @@ class InklingAttention(nn.Module): 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)} - (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 KV pool is bf16: full layers at ``out_cache_loc`` in the full pool, diff --git a/python/sglang/srt/models/inkling_common/kernels/comm.py b/python/sglang/srt/models/inkling_common/kernels/comm.py index 41dc048b6..61787c32a 100644 --- a/python/sglang/srt/models/inkling_common/kernels/comm.py +++ b/python/sglang/srt/models/inkling_common/kernels/comm.py @@ -94,7 +94,7 @@ def _ar_jit(): lazy so importing comm.py doesn't pull in the JIT machinery).""" if not is_cuda(): return None - from sglang.jit_kernel import inkling_all_reduce + from sglang.kernels.ops.model.inkling import inkling_all_reduce return inkling_all_reduce @@ -103,7 +103,7 @@ def _ar_jit(): def _ar_fused_jit(): if not is_cuda(): return None - from sglang.jit_kernel import inkling_ar_fused + from sglang.kernels.ops.model.inkling import 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 (attn-side: wo_ud AR -> attn_sconv -> mlp_norm; MoE-side: MoE AR -> 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 consuming layer/tail -- it is a pure function of per-forward state.""" if not is_cuda(): @@ -674,7 +675,7 @@ def all_gather_hidden(input: torch.Tensor, group: GroupCoordinator) -> torch.Ten def _ar_ssconv_jit(): if not is_cuda(): 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 @@ -688,7 +689,8 @@ def scattered_ar_sconv_fusable( ) -> bool: """True when an extend {reduce_scatter_hidden -> sconv(shard) -> 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 evaluate it identically.""" if not is_cuda():