fix(qsa): restore SM121 correctness with Humanize and Kernel Design Agent

python/sglang/kernels/kda_kernels/README.md

@@ -0,0 +1,10 @@
+# Kernel Design Agent kernels
+
+This directory contains optimized kernels produced through agentic kernel-design
+workflows, including [Humanize2](https://github.com/PolyArch/humanize) and
+[KDA-1.5](https://github.com/radixark/KDA-1.5).
+
+Each kernel package must document its source task, exact source revision,
+target hardware, supported shapes, and validation evidence. Generated kernels
+remain opt-in until correctness and end-to-end serving performance have been
+validated on their target GPU.

python/sglang/kernels/kda_kernels/__init__.py

@@ -0,0 +1 @@
+"""Kernels produced by Kernel Design Agent workflows."""

python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/README.md

@@ -0,0 +1,51 @@
+# Qwen3.8 QSA packed-varlen decode for SM121
+
+This implementation was optimized by Codex and Kimi K3 agents through
+[KDA-1.5](https://github.com/radixark/KDA-1.5). The task and immutable real
+tensor replay were registered in [radixark/KDA-1.5 PR #4](https://github.com/radixark/KDA-1.5/pull/4)
+at commit `414ce456e14ae8546f77d9356d2c4d955c5bb7f1`. This package integrates
+winning submission `b4181149c8884ddb`; its byte-exact submitted source has SHA256
+`4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909`.
+
+The kernel is specialized for the packed QSA decode tensors captured from
+`RadixArk/Qwen3.8-Flash-Next-NVFP4` on NVIDIA GB10 (SM121):
+
+- BF16 query, key, value, and output with head dimension 256
+- one packed query row per sequence and device-side `cu_seqlens`
+- 12 query heads per KV head: TP1 uses 24Q/2KV and TP2 uses 12Q/1KV
+- all query-row counts in the validated `1 <= bs <= 128` envelope
+- `max_seqlen_k` capacity up to 2055 and captured logical selected-KV lengths
+  up to 2051 rows per sequence
+
+The implementation groups the 12 query heads that share one KV head into one
+CTA, uses BF16 tensor-core QK/PV products with FP32 online-softmax state, and
+splits long KV rows across multiple CTAs. The last arriving split performs a
+stable FP32 merge and resets its device counter in the same launch. A
+host-visible shape/topology policy selects the two measured schedules, while
+the live device `cu_seqlens_k` selects one, two, four, or eight active splits
+without a host synchronization.
+
+SM121 dispatch checks the exact Qwen3.8 contract and routes directly to this
+kernel; it is the only packed-QSA attention implementation added by this PR.
+The KDA replay passed all 15 TP1/TP2 production tensors on two independent GB10
+GPUs, and the final source passed 150,000 consecutive launches with all
+counters returning to zero.
+
+After adaptation into SGLang, the packaged kernel passed the same 15/15 replay
+with exactly one CUDA activity per row and a 2.0702x all-shape geomean over the
+generic Triton fallback (1.6951x large, 2.3653x small). On one DGX Spark running
+the full TP1 NVFP4 model with NEXTN, three-round low-concurrency serving A/B
+improved total token throughput by 4.45% at concurrency 1 and 4.00% at
+concurrency 4. A 50-example, five-shot GSM8K A/B with a 2048-token output limit
+scored 49/50 for both Triton and KDA, with the same single failed example.
+
+An additional synthetic GB10 sweep covers both TP topologies, every batch size
+from 1 through 16, and short plus saturated KV rows. All 64 cases passed; the
+maximum relative L2 against the original correct Triton implementation was
+0.002422, and speedup ranged from 1.41x to 5.09x.
+
+A follow-up extended-batch sweep covers both TP topologies, batch sizes
+17/24/32/48/64/96/128, and short, medium, plus saturated KV rows. All 42 cases
+passed with maximum relative L2 0.002410. Geomean speedup was 4.48x, the slowest
+case still improved by 1.58x, and no case regressed. The packaged scratch space
+is therefore sized for the largest tested batch, 128.

python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/__init__.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# KDA provenance: optimized by Codex and Kimi K3 agents through KDA-1.5.
+# Task: https://github.com/radixark/KDA-1.5/pull/4 @
+# 414ce456e14ae8546f77d9356d2c4d955c5bb7f1.
+# Winning submission: b4181149c8884ddb.
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_HEAD_TOPOLOGIES = frozenset({(12, 1), (24, 2)})
+# Largest batch qualified by the extended GB10 baseline sweep.
+_MAX_BATCH = 128
+_MAX_SELECTED_KV = 2055
+_logged_fast_path = False
+
+
+def can_use_qwen38_qsa_sm121(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+) -> bool:
+    """Return whether this call matches the captured Qwen3.8 SM121 contract."""
+    if not q.is_cuda or q.ndim != 3 or q.dtype != torch.bfloat16:
+        return False
+    batch, num_q_heads, head_dim = q.shape
+    if not (0 < batch <= _MAX_BATCH) or head_dim != 256:
+        return False
+    if k.ndim != 3 or v.shape != k.shape or k.dtype != q.dtype or v.dtype != q.dtype:
+        return False
+    num_kv_heads = k.shape[1]
+    if (num_q_heads, num_kv_heads) not in _SUPPORTED_HEAD_TOPOLOGIES:
+        return False
+    if k.shape[2] != head_dim or not (0 < max_seqlen_k <= _MAX_SELECTED_KV):
+        return False
+    if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
+        return False
+    if q.device != k.device or q.device != v.device:
+        return False
+    if cu_seqlens_q.device != q.device or cu_seqlens_k.device != q.device:
+        return False
+    if cu_seqlens_q.dtype != torch.int32 or cu_seqlens_k.dtype != torch.int32:
+        return False
+    if cu_seqlens_q.ndim != 1 or cu_seqlens_k.ndim != 1:
+        return False
+    if not cu_seqlens_q.is_contiguous() or not cu_seqlens_k.is_contiguous():
+        return False
+    if cu_seqlens_q.numel() != batch + 1 or cu_seqlens_k.numel() != batch + 1:
+        return False
+    properties = torch.cuda.get_device_properties(q.device)
+    return (properties.major, properties.minor) == (12, 1)
+
+
+def qwen38_qsa_sm121(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Run the KDA-generated Qwen3.8 packed QSA decode kernel."""
+    global _logged_fast_path
+    if not can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k):
+        raise ValueError("unsupported call for the KDA Qwen3.8 SM121 QSA kernel")
+
+    from .kernel import qwen38_qsa_sm121 as run_kernel
+
+    if not _logged_fast_path:
+        logger.info(
+            "Using the Codex/Kimi K3 KDA Qwen3.8 QSA kernel on SM121 "
+            "(radixark/KDA-1.5#4, submission b4181149c8884ddb)"
+        )
+        _logged_fast_path = True
+    return run_kernel(q, k, v, cu_seqlens_q, cu_seqlens_k, softmax_scale)
+
+
+__all__ = ["can_use_qwen38_qsa_sm121", "qwen38_qsa_sm121"]