[None][chore] Reconcile K3 KDA/attn-res kernel files with main (PR #17225) - #17230
[None][chore] Reconcile K3 KDA/attn-res kernel files with main (PR #17225)#17230brnguyen2 wants to merge 2 commits into
Conversation
…IDIA#17225) PR NVIDIA#17225 carries these kernels to main with review fixes and repo-wide clang-format/ruff formatting that this branch predates. Sync the file contents so the eventual mergeback of feat/kimi_k3 into main does not conflict on them, and so this branch picks up the two functional fixes from that review: - fused_k123: zero-initialize all 64 SMEM rows (rows 62-63 previously read stale SMEM through the sAkk transposed store in the non-pure varlen path) - attnResFwd.cu: add the missing __syncwarp before the tail cross-lane logits reads Everything else is formatting and comment updates. Feat-side changes made after the kernels were exported (per-file 3-way merge against the export base) are preserved; no feat-only functionality is reverted. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #63578 [ run ] triggered by Bot. Commit: |
|
PR_Github #63578 [ run ] completed with state
|
The review-round change narrowing the k4_persistent monkey-patch handler
to `except (AttributeError, ImportError)` breaks every cute.compile of
the module: the nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse
tuple except handlers ("'Tuple' object has no attribute 'id'"). Revert
to a single bare Exception with a comment explaining the constraint.
Also gate the new A_log validation on a non-empty token batch: zero-token
calls take the early return in _chunk_kda_fwd and never touch A_log, and
the runtime emits such batches under the overlap scheduler + logprobs
flows.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #63751 [ run ] triggered by Bot. Commit: |
|
Pre-merge validation results for this PR at the current tip (42f4f06), run on a self-built checkout of this branch on a Blackwell (sm_100 family) node: Kernel-vs-reference parity unit suites — all green:
Logits-parity integration run — 4-GPU TP4 generation on a layer-truncated checkpoint: the baseline leg and the spec-dec logits-parity leg both pass (52 prompts, parity verified, zero drift). Two findings on the target branch itself, orthogonal to this PR's content:
Validation history: an earlier run at 710b0a4 caught that narrowing a module-level handler in k4_persistent.py to |
|
PR_Github #63751 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63765 [ run ] triggered by Bot. Commit: |
|
PR_Github #63765 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63779 [ run ] triggered by Bot. Commit: |
|
PR_Github #63779 [ run ] completed with state
|
2 participants
Description
#17225 carries the KDA prefill/MTP decode CuTe DSL kernels and the fused attention-residual kernel to main, with review fixes and repo-standard clang-format/ruff formatting that this branch predates. This PR syncs the file contents on
feat/kimi_k3to match, so that:fused_k123: zero-initialize all 64 SMEM rows — rows 62–63 previously read stale shared memory through thesAkktransposed store in the non-pure varlen path;attnResFwd.cu: add the missing__syncwarpbefore the tail cross-lane logits reads.Everything else is formatting and comment updates. Feat-side changes made after the kernels were exported were preserved via per-file three-way merges against the export base; no feat-only functionality is reverted. CMakeLists and
custom_ops/__init__.pyare untouched (they already agree or carry feat-only registrations).Test Coverage
The kernel content matches #17225's tip, which passed a per-architecture compile check (sm_80/90/100f build; sm_120f rejected as designed by the arch pin). The KDA kernel parity unit suites on this branch cover the merged files; a parity rerun over the review's device-code changes is queued and should gate the merge of this PR.
PR Checklist