refactor(engine/hip): kernels take the QuixiCore shape — a family per folder - #22
Conversation
… family
rocm_kernels.hip was one 17,090-line mega file. It is now a 72-line
amalgamation translation unit: preamble, the anonymous namespace, and an
ordered #include manifest over 59 fragments in 13 semantic families —
contract, common, activations, attention, embedding, lora, matmul, moe,
norms, quant, sampling, serving, training. The taxonomy is QuixiCore's
(kernels/<family>/, contract central, implementations family-local),
adapted: no variants/<target>/ tree, because this engine compiles ONE
source for AMD, NVIDIA and HIP-CPU alike, so the variant axis is empty.
Pure reorganisation — no maths changed, no kernel symbol renamed. Every
fragment is a contiguous slice of the original TU, so the manifest
reassembles it exactly.
Byte-identity receipts:
- all 142 exported kernel bodies hash-identical before -> after
(sha256 of the 142-row table: 38680f7805edeffbb5028bb8fda9b0c1d0d08b98cb037864cd5f91a71709f3e4)
- reassembled source == HEAD's rocm_kernels.hip (454c04d0...) except
the manifest's own two `// ---` phase banners; nothing else differs
- the anonymous namespace (lines 8..5092 of the original, four KV
kernels inside it) is preserved verbatim — ordering is load-bearing
in a single TU and the manifest phases encode it
The manifest is explicit, never a glob: a fragment on disk that nobody
wired would compile to nothing. TestHIPKernelSource_FragmentManifestCovers
EveryFile_Good fails on an unwired or twice-wired fragment (verified
against a planted orphan) and on a missing licence header. The 34 source
assertions read through hipReadKernelSource, which resolves the manifest,
so every existing check runs against the same assembled text.
Receipts — Mac (darwin/arm64):
GOWORK=off go build ./... exit 0
GOWORK=off go vet ./... exit 0
task test FAIL, pre-existing: engine/metal
TestLaneSetThroughputAB + TestLaneSetGEMME2BByteIdentityHiddens fail
on a symlinked gemma4e2b/config.json ("path escapes from parent").
Same FAIL at 2727995 before this change; this lane touches no file
under engine/metal.
Receipts — homelab (linux/amd64, ROCm 7.2, gfx1101, CUDA 12.8):
make hip-amd AMD_HIP_ARCH=gfx1101 exit 0 1386800 B (= baseline)
make hip-nvidia CUDA_PATH=cuda-12.8 exit 0 5992208 B (= baseline)
make hip-cpu-x86_64 exit 0 33576000 B (= baseline)
go build ./... exit 0
go vet ./engine/hip exit 0
go test ./engine/hip -count=1 ok 4.587s
42 TestHIPKernelSource_* pass over the new layout
Object hashes are NOT compared: hipcc is not byte-deterministic here —
two compiles of one unchanged source differ in 904 bytes. Output sizes
match on all three lanes; the identity claim rests on the source.
Co-Authored-By: Virgil <virgil@lethean.io>
📝 WalkthroughWalkthroughThe PR adds a manifest-based HIP kernel source assembly path and documentation. It introduces a shared launch ABI, device validators, and ROCm kernels for attention, quantisation, serving, sampling, embeddings, normalisation, routing, activations, projections, and training losses. Existing source tests now inspect assembled fragments. ChangesROCm HIP kernel platform
Sequence Diagram(s)sequenceDiagram
participant LaunchPacket
participant DeviceValidator
participant HIPKernel
participant DeviceBuffers
LaunchPacket->>DeviceValidator: provide ABI metadata and buffer descriptors
DeviceValidator->>DeviceValidator: validate versions, sizes, pointers, and dimensions
DeviceValidator-->>HIPKernel: allow valid launch
HIPKernel->>DeviceBuffers: read encoded inputs and write outputs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)Error: can't load config: unsupported version of the configuration: "" See https://golangci-lint.run/docs/product/migration-guide for migration instructions Warning Billing warning: we have not been able to collect payment for this subscription for more than 72 hours. Please update the payment method or pay any pending invoices in Billing to avoid service interruption. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
go/engine/hip/kernels/README.md (1)
101-104: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCorrect the C++ standard requirement.
The NVIDIA and HIP-CPU compile paths use C++20. The statement that the HIP source is built as C++23 can cause contributors to add code that breaks supported backends.
State that AMD builds use C++23 and that shared source must remain C++20-compatible.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/README.md` around lines 101 - 104, Update the build-requirement statement in the HIP kernels README to say that AMD builds use C++23 while shared source must remain compatible with C++20 for NVIDIA and HIP-CPU compile paths. Remove the claim that HIP source generally uses C++23, while preserving the existing cgo and mdspan guidance.
🟠 Major comments (24)
go/engine/hip/kernels/training/losses.hipinc-91-94 (1)
91-94: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject non-finite rewards before device dispatch.
hipLoadedModel.RunGRPOAdvantageaccepts any non-empty reward slice atgo/engine/hip/hip_training_launch.goLines 732-747. ANaNor infinity reaches this guard. The guard returns without writingoutput. The launcher then reads the output after a successful kernel launch, so the caller receives no error for an invalid advantage result.Validate every reward on the host and return an error before
hipRunGRPOAdvantageKernel.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/training/losses.hipinc` around lines 91 - 94, Validate every element of the reward slice in hipLoadedModel.RunGRPOAdvantage before invoking hipRunGRPOAdvantageKernel, rejecting NaN and infinite values with an error. Keep the existing non-empty input checks and ensure invalid rewards return before device dispatch or output access.go/engine/hip/kernels/contract/launch_abi.hipinc-1848-1850 (1)
1848-1850: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAdd the missing packed top-k ABI assertion.
rocm_packed_topk_launch_argshas nostatic_assertagainstROCM_PACKED_TOPK_LAUNCH_ARGS_BYTES. A field or alignment change can then compile while the launcher packet contract has drifted.Proposed fix
+static_assert(sizeof(rocm_packed_topk_launch_args) == ROCM_PACKED_TOPK_LAUNCH_ARGS_BYTES, "packed top-k launch ABI drift"); static_assert(sizeof(rocm_packed_topk_sample_launch_args) == ROCM_PACKED_TOPK_SAMPLE_LAUNCH_ARGS_BYTES, "packed top-k sample launch ABI drift");🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/contract/launch_abi.hipinc` around lines 1848 - 1850, Add a static_assert for rocm_packed_topk_launch_args in the ABI assertion section, comparing sizeof(rocm_packed_topk_launch_args) with ROCM_PACKED_TOPK_LAUNCH_ARGS_BYTES and using a clear packed top-k launch ABI drift message. Leave the existing assertions unchanged.go/engine/hip/kernels/attention/single_head.hipinc-836-841 (1)
836-841: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winRemove the unused per-token page lookup in the fallback loop.
Line 838 assigns
pageand never reads it. Line 839 callsrocm_attention_device_kv_value, which resolves the page again throughrocm_attention_device_kv_page.This loop is nested inside the
dimloop at Line 834. For the device-KV path the code therefore performsargs.dim * args.token_countredundant descriptor lookups, and each lookup can run a binary search overpage_count. Delete the assignment.Note: this fragment is a move of pre-existing code, so this is not a regression introduced by the PR.
⚡ Proposed fix
for (uint32_t dim = tid; dim < args.dim; dim += threads) { float out = 0.0f; for (uint32_t token = 0; token < args.token_count; ++token) { const uint32_t base = token * kv_stride; - const rocm_device_kv_page_descriptor *page = device_kv ? rocm_attention_device_kv_page(args, token) : nullptr; const float value = device_kv ? rocm_attention_device_kv_value(args, false, token, dim) : values[base + dim]; out += weights[token] * value; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/attention/single_head.hipinc` around lines 836 - 841, Remove the unused page assignment calling rocm_attention_device_kv_page from the per-token fallback loop around rocm_attention_device_kv_value, while preserving the base, value, and accumulation logic.go/engine/hip/kernels/attention/chunked.hipinc-96-110 (1)
96-110: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate page geometry on the direct token-page fast path.
When
direct_kq8vq4_token_pagesis true, Line 97 reduces validation to three null checks. The slow path at Line 99 instead callsrocm_attention_kq8vq4_page_valid(page, token, kv_width), which checkstoken_count,key_width == kv_width,value_width == kv_width, the q8 and q4 encodings, and the declaredkey_bytesandvalue_bytes.The fast path then uses the unchecked fields. Line 105 adds
kv_dim_offsetto the key payload pointer, Line 109 adds(value_base + kv_dim_offset) >> 1uto the value payload pointer, and Lines 117-120 readkey_values[dim]for everydim < args.dim. If a page reportskey_widthsmaller thankv_width = args.key_heads * args.dim, those reads run past the page allocation. The headermode_code == ROCM_DEVICE_KV_DESCRIPTOR_MODE_KQ8VQ4check at Line 71 describes the stream, not each page.The same asymmetry repeats at Lines 147-151, 441-445, 492-496, 730-734, 793-797, 1103-1107, 1165-1167, 1447-1451 and 1525-1527.
Either check the widths and encodings on the fast path, or hoist a one-time geometry check on the first page before
direct_kq8vq4_token_pagesis set.Note: this fragment is a move of pre-existing code, so this is not a regression introduced by the PR.
🛡️ Proposed fix for the fast-path check
if (direct_kq8vq4_token_pages) { - page_valid = page != nullptr && page->key_pointer != 0 && page->value_pointer != 0 ? 1u : 0u; + page_valid = page != nullptr && + page->key_pointer != 0 && + page->value_pointer != 0 && + page->key_width == kv_width && + page->value_width == kv_width && + rocm_device_kv_encoding_is_q8(page->key_encoding) && + rocm_device_kv_encoding_is_q4(page->value_encoding) ? 1u : 0u; } else { page_valid = rocm_attention_kq8vq4_page_valid(page, token, kv_width) ? 1u : 0u; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/attention/chunked.hipinc` around lines 96 - 110, Validate full page geometry before using the direct token-page fast path: update the direct_kq8vq4_token_pages handling around rocm_attention_kq8vq4_page_valid so it checks token_count, key_width/value_width against kv_width, required q8/q4 encodings, and declared key_bytes/value_bytes, or perform an equivalent one-time validation before enabling the fast path. Apply the same protection to every repeated direct-page validation site in this file, including the blocks near the referenced attention paths, while preserving the existing slow-path validation.go/engine/hip/kernels/attention/device_kv.hipinc-250-251 (1)
250-251: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard the packed
int32_tload against a misaligned payload pointer.
row_base_indexonly checks the row offset withinpayload. For row-interleaved q8,payloadcan be misaligned whenvaluesis 4-byte aligned but width makeslocal_token * stride + sizeof(uint32_t)non-4-byte aligned. In that fast path, cast throughuintptr_tbefore the grouped q8 load or skip it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/attention/device_kv.hipinc` around lines 250 - 251, Update the packed q8 fast path around the row_base_index check to also validate that values + row_base_index is 4-byte aligned before casting to const int32_t*. Use a uintptr_t-based alignment check, and fall back to the existing safe load path when misaligned.go/engine/hip/kernels/attention/chunked.hipinc-1042-1050 (1)
1042-1050: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRestrict
chunk_sizeto the per-variant shared-memory budget.
rocm_attention_heads_batch_chunked_stage1_gqa4andgqa8allocate per-block shared buffers proportional togroup_size * chunk_size * dim, but the common config validation only rejects tiny/invalidchunk_size; it does not check the GQA shared-memory size before the selected kernel launch. Add the shared-memory calculation to the selected GQA path or reject configurations whose calculatedSharedMemBytesexceeds the chosen kernel’s limit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/attention/chunked.hipinc` around lines 1042 - 1050, Add per-variant shared-memory validation for the selected GQA path before launching rocm_attention_heads_batch_chunked_stage1_gqa4 or rocm_attention_heads_batch_chunked_stage1_gqa8. Calculate SharedMemBytes using the selected group_size, chunk_size, and dim, then reject configurations exceeding that kernel’s shared-memory limit while preserving existing validation for valid configurations.go/engine/hip/kernels/lora/projection.hipinc-25-41 (1)
25-41: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftHoist the LoRA down-projection out of the per-row loop.
The inner loop at lines 35-37 computes
down = sum_col lora_a[r][col] * input[col]. This value depends only onrandinput. It does not depend onrow. Every thread in the grid therefore recomputes the samerank * colsreduction.Total work becomes
rows * rank * colsinstead ofrank * cols. For 4096 rows and rank 16, the LoRA term costs about 4096 times more than required, and it dominates the kernel.The access pattern is also uncoalesced. Consecutive threads read
base[row * args.cols + col]at a stride ofargs.cols.rocm_projectioningo/engine/hip/kernels/matmul/projection.hipincuses one block per row withrocm_block_reduce_sum, which reads contiguously.Compute
down[r]once per block into shared memory, then reuse it for every row. Alternatively, adopt the block-per-row layout used byrocm_projection.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/lora/projection.hipinc` around lines 25 - 41, Hoist the LoRA down-projection out of the per-row computation in the projection kernel: compute each rank component of down once per block, store the results in shared memory, synchronize, and reuse them while accumulating delta for each row. Update the kernel’s launch/shared-memory handling as needed, preserving the existing bias, scaling, and output behavior.go/engine/hip/kernels/norms/rms_norm_validate.hipinc-146-147 (1)
146-147: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick win32-bit dimension products wrap before widening, in both validators and index arithmetic. Each site multiplies two
uint32_tdimension fields. C++ evaluates the product in 32-bit arithmetic and wraps modulo 2^32. Widening bysizeof(float)orsizeof(uint16_t)happens only afterwards, so it does not prevent the wrap. The result is a byte check that accepts an undersized buffer, or a device index that addresses the wrong memory. Several validators in the same files already use the correct 64-bit form, which makes the fix unambiguous.
go/engine/hip/kernels/norms/rms_norm_validate.hipinc#L146-L147: computehead_dim * head_countasuint64_t, bound it by0xffffffffu, and use it in both byte comparisons. Apply the same change torocm_valid_rms_norm_rope_heads_argsat lines 177-178.go/engine/hip/kernels/lora/projection_validate.hipinc#L22-L24: castargs.rows,args.rankandargs.colstouint64_tin thebase_weight_bytes,lora_a_bytesandlora_b_bytescomparisons.go/engine/hip/kernels/matmul/projection.hipinc#L29-L29: castrowtouint64_tinweights[row * args.cols + col], and repeat at lines 34, 40, 45, 86, 98, 111 and 123.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/norms/rms_norm_validate.hipinc` around lines 146 - 147, Prevent 32-bit dimension-product overflow across the validators and index arithmetic. In go/engine/hip/kernels/norms/rms_norm_validate.hipinc lines 146-147 and 177-178, compute head_dim * head_count as uint64_t, bound it by 0xffffffffu, and reuse it for both byte checks. In go/engine/hip/kernels/lora/projection_validate.hipinc lines 22-24, cast rows, rank, and cols to uint64_t in the base_weight_bytes, lora_a_bytes, and lora_b_bytes comparisons. In go/engine/hip/kernels/matmul/projection.hipinc lines 29, 34, 40, 45, 86, 98, 111, and 123, cast row to uint64_t before multiplying by args.cols in each weights index.go/engine/hip/kernels/norms/rms_norm_fused_projection.hipinc-119-119 (1)
119-119: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAlign the
output_scale_bits == 0convention with the other RMS kernels.Line 119 converts
args.output_scale_bitsdirectly. A value of 0 therefore produces a scale of 0.0f and zeroes bothresidual_outputand the projection input.The sibling kernels treat 0 as "unset" and substitute 1.0f. See
go/engine/hip/kernels/norms/rms_norm.hipincline 97 and line 200, and the matching validators atgo/engine/hip/kernels/norms/rms_norm_validate.hipincline 79 and line 110. The validator for this kernel only checksisfinite(output_scale), so 0 passes validation silently.If the host always sets a non-zero scale for this kernel, reject 0 in the validator instead. If the host may omit the scale, apply the same substitution here.
🔧 Proposed fix: use the shared convention
- const float output_scale = rocm_float_from_bits(args.output_scale_bits); + const float output_scale = args.output_scale_bits == 0 ? 1.0f : rocm_float_from_bits(args.output_scale_bits);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/norms/rms_norm_fused_projection.hipinc` at line 119, Update the output-scale handling in the fused projection kernel around output_scale and its validator so output_scale_bits == 0 follows the shared RMS-kernel convention and resolves to 1.0f before scaling residual_output or the projection input. Reuse the existing bit-decoding path and preserve nonzero scale behavior; only reject zero in the validator if this kernel’s host contract guarantees the scale is always provided.go/engine/hip/kernels/matmul/projection_validate.hipinc-5-16 (1)
5-16: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winValidate the bias pointer and bias size in
rocm_valid_projection_args.
rocm_valid_projection_argsnever inspectsargs.bias_pointer,args.bias_bytesorROCM_PROJECTION_LAUNCH_FLAG_BIAS.rocm_projectionreadsbias[row]atgo/engine/hip/kernels/matmul/projection.hipincline 52 for anyrowbelowargs.rows. A bias buffer shorter thanrows * sizeof(float)produces an out-of-bounds device read.
rocm_valid_projection_batch_argsin this file already performs both checks at lines 46-53. Apply the same checks to the scalar validator.🛡️ Proposed fix
args.output_bytes != args.rows * sizeof(float)) { return false; } + if ((args.flags & ROCM_PROJECTION_LAUNCH_FLAG_BIAS) != 0 && + (args.bias_pointer == 0 || args.bias_bytes != args.rows * sizeof(float))) { + return false; + } + if ((args.flags & ROCM_PROJECTION_LAUNCH_FLAG_BIAS) == 0 && + (args.bias_pointer != 0 || args.bias_bytes != 0)) { + return false; + } if (args.weight_encoding == ROCM_PROJECTION_WEIGHT_ENCODING_FP16 ||🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/matmul/projection_validate.hipinc` around lines 5 - 16, Update rocm_valid_projection_args to validate the bias configuration like rocm_valid_projection_batch_args: require the bias pointer when ROCM_PROJECTION_LAUNCH_FLAG_BIAS is set, and require bias_bytes to equal rows * sizeof(float) when bias is enabled. Preserve validation for configurations without bias.go/engine/hip/kernels/embedding/mean_pool_validate.hipinc-11-12 (1)
11-12: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winByte-count checks in the embedding validators use 32-bit products. Both validators multiply two
uint32_tcount fields together before thesizeof(float)operand promotes the expression tosize_t. The 32-bit product can wrap, so a malformed packet can pass a byte-count check that does not cover the buffer, and the kernel then reads out of bounds. The MoE validators ingo/engine/hip/kernels/moe/routing_validate.hipincalready apply a leadingstatic_cast<uint64_t>at Lines 41, 42, 43, 55, and 56. Apply that same pattern here.
go/engine/hip/kernels/embedding/mean_pool_validate.hipinc#L11-L12: castargs.token_counttouint64_tbefore multiplying byargs.dim, and castargs.dimtouint64_tin the output check.go/engine/hip/kernels/embedding/rerank_validate.hipinc#L13-L13: castargs.document_counttouint64_tbefore multiplying byargs.dim.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/embedding/mean_pool_validate.hipinc` around lines 11 - 12, Prevent 32-bit overflow in the embedding byte-count validators by promoting operands before multiplication. In go/engine/hip/kernels/embedding/mean_pool_validate.hipinc lines 11-12, cast args.token_count before multiplying by args.dim and cast args.dim in the output-size check; in go/engine/hip/kernels/embedding/rerank_validate.hipinc line 13, cast args.document_count before multiplying by args.dim. Use the existing uint64_t promotion pattern from the MoE validators.go/engine/hip/kernels/moe/routing_validate.hipinc-80-82 (1)
80-82: 🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy liftAdd size checks for the chunk-held MLX affine buffers.
rocm_moe_mlx_affine_routes_chunkonly stores pointers forgate_up_*anddown_*buffers, and validation rejects zero pointers but reads pastargs.rows * packed_per_rowandargs.cols / args.group_size. Model the kernel-expected extents for each supported bit width and validate them at launch to avoid trusting undersized chunks that pass pointer checks.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/moe/routing_validate.hipinc` around lines 80 - 82, Update the validation logic in rocm_moe_mlx_affine_routes validation to verify the byte sizes of every chunk-held gate_up_* and down_* buffer, not just pointer validity. Compute the expected extents from args.rows, args.cols, args.group_size, and the supported bit width, including packed_per_row, and reject undersized buffers before launch while preserving the existing input, chunk, and output checks.go/engine/hip/kernels/quant/affine_projection_validate.hipinc-46-47 (1)
46-47: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSize checks in this layer multiply two 32-bit fields before the widening operand. Each affected comparison has the form
count_a * count_b * sizeof(T). C++ evaluates left to right, socount_a * count_bstays 32-bit and wraps at 2^32. Only then doessizeof(T)widen the result. A packet whose product wraps therefore passes the byte check against a much smaller allocation, and the kernels index past that allocation.go/engine/hip/kernels/sampling/diffusion_validate.hipinclines 14-15 show the correct pattern for this layer.
go/engine/hip/kernels/quant/affine_projection_validate.hipinc#L46-L47: applystatic_cast<uint64_t>(args.batch)to theinput_bytesandoutput_bytesproducts inrocm_valid_mlx_q4_projection_batch_args.go/engine/hip/kernels/quant/jangtq_codebook_validate.hipinc#L60-L61: applystatic_cast<uint64_t>toargs.codebook_countandargs.code_countinrocm_valid_codebook_args.go/engine/hip/kernels/quant/affine_projection_validate.hipinc#L88-L89: applystatic_cast<uint64_t>(args.batch)to both products inrocm_valid_mlx_q4_projection_greedy_batch_args.Note the PR states the reassembled source must match the original except for two banner comments. If you accept these fixes, update the source-comparison expectation in
go/engine/hip/hip_kernel_source_test.go.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/quant/affine_projection_validate.hipinc` around lines 46 - 47, Prevent 32-bit multiplication overflow in the validation byte checks by widening the first multiplicand before each product. In go/engine/hip/kernels/quant/affine_projection_validate.hipinc:46-47, update both products in rocm_valid_mlx_q4_projection_batch_args; in go/engine/hip/kernels/quant/jangtq_codebook_validate.hipinc:60-61, widen args.codebook_count and args.code_count in rocm_valid_codebook_args; and in go/engine/hip/kernels/quant/affine_projection_validate.hipinc:88-89, update both products in rocm_valid_mlx_q4_projection_greedy_batch_args. Also update the source-comparison expectation in go/engine/hip/hip_kernel_source_test.go to account for these accepted source changes.go/engine/hip/kernels/quant/jangtq_codebook.hipinc-25-29 (1)
25-29: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winWiden the packed element index to 64-bit.
Line 26 computes
row * args.colsin 32-bit arithmetic.rocm_unpack_signed_bitsthen computesindex * bits, also in 32-bit, at line 32 ofgo/engine/hip/kernels/quant/jangtq_codebook_validate.hipinc.The bit offset wraps first. For
bits == 8the wrap starts oncerows * colsexceeds about 5.4e8 elements, which is roughly 512 MiB of packed weights. After the wrap the kernel reads a valid but wrong byte, so the projection produces silently incorrect results rather than faulting.The validator already widens for the same quantity. Line 6 of
jangtq_codebook_validate.hipincusesstatic_cast<uint64_t>(args.bits) * args.rows * args.cols. The kernel and the validator therefore disagree on the addressable range.🐛 Proposed fix for the index arithmetic
Change the helper to take a 64-bit index:
-__device__ int8_t rocm_unpack_signed_bits(const uint8_t *packed, uint32_t bits, uint32_t index) +__device__ int8_t rocm_unpack_signed_bits(const uint8_t *packed, uint32_t bits, uint64_t index) { - const uint32_t bit_offset = index * bits; - const uint32_t byte_index = bit_offset / 8u; + const uint64_t bit_offset = index * bits; + const uint64_t byte_index = bit_offset / 8u; const uint32_t shift = bit_offset % 8u;Then widen the caller:
for (uint32_t col = 0; col < args.cols; ++col) { - const uint32_t index = row * args.cols + col; + const uint64_t index = static_cast<uint64_t>(row) * args.cols + col; const int8_t quantized = rocm_unpack_signed_bits(packed, args.bits, index); sum += input[col] * static_cast<float>(quantized) * scale; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/quant/jangtq_codebook.hipinc` around lines 25 - 29, Widen packed-element index arithmetic in the quantization kernel: update rocm_unpack_signed_bits and its callers to accept and compute with uint64_t, and cast before multiplying row by args.cols in the kernel loop. Ensure the subsequent index * bits calculation also remains 64-bit, matching the validator’s addressable range.go/engine/hip/kernels/sampling/diffusion.hipinc-501-505 (1)
501-505: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe output store has no per-dimension bounds check.
Line 503 writes
dim_base + dim_lanefor all four values. The only guard isdim_valid, set at Line 456 fromdim_base < args.hidden_size. Ifargs.hidden_sizeis not a multiple ofROCM_DIFFUSION_EXPECTED_EMBEDDING_Q8_G64_DIMS_PER_THREAD, the final thread writes up to three elements past the end of the row. On the last row that write leaves the output allocation.
rocm_diffusion_expected_embedding_q8_g64_dims4_rows4performs the same tiled store and does check each dimension at Line 418. The same gap exists at Line 593 inrocm_diffusion_expected_embedding_q8_g64_subgroup32_rows64_prob4.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/sampling/diffusion.hipinc` around lines 501 - 505, Add a per-dimension bounds check to the output stores in the affected kernels, including rocm_diffusion_expected_embedding_q8_g64_dims4_rows4 and rocm_diffusion_expected_embedding_q8_g64_subgroup32_rows64_prob4, so each write occurs only when dim_base + dim_lane is less than args.hidden_size. Preserve the existing accumulation and scaling behavior for valid dimensions.go/engine/hip/kernels/quant/autoround.hipinc-19-22 (1)
19-22: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
packed_start_bitoverflows 32 bits before the largest validated tensor.
value_startandpacked_start_bitareuint32_t.rocm_valid_autoround_quantize_argscomputespacked_bitsasuint64_t, so it accepts tensors whose total bit offset exceedsUINT32_MAX. Atbits == 8, that limit is reached at about 536 million weights, which is roughly 2 GiB offloatinput. The validator accepts such a launch.packed_start_bytethen wraps and the kernel writes to the wrong offset inpacked.Widen the offsets to 64-bit.
rocm_autoround_pack_signedneeds the same treatment for itsindexparameter.🐛 Proposed fix: widen the offsets
- const uint32_t value_start = row * args.cols + group * args.group_size; - const uint32_t packed_start_bit = value_start * args.bits; - const uint32_t packed_start_byte = packed_start_bit >> 3; + const uint64_t value_start = static_cast<uint64_t>(row) * args.cols + static_cast<uint64_t>(group) * args.group_size; + const uint64_t packed_start_bit = value_start * args.bits; + const uint64_t packed_start_byte = packed_start_bit >> 3; const uint32_t packed_group_bytes = (args.group_size * args.bits + 7u) / 8u;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/quant/autoround.hipinc` around lines 19 - 22, Widen the offset arithmetic in the autoround quantization kernel to 64-bit: update value_start, packed_start_bit, packed_start_byte, and packed_group_bytes in the shown launch-offset calculation, and change the index parameter of rocm_autoround_pack_signed to the matching 64-bit type. Preserve the existing offset formulas while ensuring multiplication, shifting, and indexing cannot overflow before addressing packed data.go/engine/hip/kernels/sampling/greedy.hipinc-39-75 (1)
39-75: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe tree reduction drops candidates when
blockDim.xis not a power of two.Line 39 accepts any
threadsvalue in the range 1 to 256. The reduction at Line 65 starts atstride = threads >> 1and halves. That schedule only covers every slot whenthreadsis a power of two.Take
threads == 6. The first pass usesstride == 3and merges slots 3, 4, 5 into slots 0, 1, 2. The second pass usesstride == 1and merges slot 1 into slot 0. Slot 2 is never merged. The candidate held by thread 2 is lost, androcm_softcap_greedy_samplecan return a token that is not the argmax.Two fixes exist. Reject a non-power-of-two
threadsat Line 39, or round the first stride up to the next power of two and bound the partner index.🐛 Proposed fix: handle a non-power-of-two thread count
- for (uint32_t stride = threads >> 1; stride > 0; stride >>= 1) { - if (tid < stride) { + uint32_t reduction_width = 1u; + while (reduction_width < threads) { + reduction_width <<= 1u; + } + for (uint32_t stride = reduction_width >> 1; stride > 0; stride >>= 1) { + if (tid < stride && tid + stride < threads) { const float other_score = scores[tid + stride]; const int32_t other_index = indices[tid + stride]; if (other_index >= 0 && (indices[tid] < 0 || other_score > scores[tid] || (other_score == scores[tid] && other_index < indices[tid]))) { scores[tid] = other_score; indices[tid] = other_index; } } __syncthreads(); }The alternative is a stricter launch check:
- if (threads == 0 || threads > 256) { + if (threads == 0 || threads > 256 || (threads & (threads - 1u)) != 0u) { return; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/sampling/greedy.hipinc` around lines 39 - 75, Update the tree reduction in rocm_softcap_greedy_sample so non-power-of-two threads are handled correctly. Either reject non-power-of-two values in the existing threads validation, or start reduction at the next power of two and guard partner indices before reading scores and indices; preserve the argmax behavior for all accepted thread counts.go/engine/hip/kernels/quant/autoround.hipinc-26-33 (1)
26-33: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winThe non-finite early return leaves the group output uninitialised.
If any weight in the group is not finite, the thread returns at Line 30. The thread then skips the scale store at Line 37, the byte clear at Lines 38-40, and the packing loop. The destination buffers keep their previous contents. The host receives a partially written result and has no status channel to detect it.
Two options exist. Write a deterministic sentinel for the group, or reject non-finite weights during validation. A deterministic sentinel keeps the kernel total.
🛡️ Proposed fix: write a deterministic zeroed group
float max_abs = 0.0f; + bool finite = true; for (uint32_t offset = 0; offset < args.group_size; ++offset) { const float value = weights[value_start + offset]; if (!isfinite(value)) { - return; + finite = false; + break; } max_abs = fmaxf(max_abs, fabsf(value)); } + if (!finite) { + scales[group_index] = 0.0f; + for (uint32_t byte = 0; byte < packed_group_bytes; ++byte) { + packed[packed_start_byte + byte] = 0; + } + return; + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/quant/autoround.hipinc` around lines 26 - 33, Update the non-finite handling in the group-processing kernel around the max_abs scan so it does not return before writing outputs. For any group containing a non-finite weight, write a deterministic zero scale, clear the group’s destination bytes, and skip packing while preserving normal processing for finite groups.go/engine/hip/kernels/serving/tiny_validate.hipinc-5-15 (1)
5-15: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win32-bit size products in the tiny validators allow out-of-bounds access in the tiny kernels. Both tiny validators compute element counts as
uint32_t * uint32_tbefore promoting tosize_t. The product wraps at 2^32, so the validator accepts a buffer that is far smaller than the kernel then addresses.rocm_valid_tiny_decode_argsalso computes(args.prior_token_count + 1u)in 32 bits, which wraps to zero atUINT32_MAX.
go/engine/hip/kernels/serving/tiny_validate.hipinc#L5-L15: cast one operand touint64_tin thetable_countproduct at line 5, and apply the same widening to the byte checks at lines 30, 35 and 36, and to lines 41, 67, 68, 69, 72, 73 and 74 inrocm_valid_tiny_decode_args.go/engine/hip/kernels/serving/tiny.hipinc#L106-L119: no change is needed here once the validator computes the sizes in 64 bits. Confirm after the fix thatupdated_baseat line 111 and the copy bound at line 107 stay inside the validatedupdated_key_bytesandupdated_value_bytes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/serving/tiny_validate.hipinc` around lines 5 - 15, Widen all tiny validator size arithmetic to 64 bits before multiplication or addition: update table_count and the byte checks in rocm_valid_tiny_decode_args, including the prior_token_count plus one calculation, by casting an operand to uint64_t. In tiny.hipinc at lines 106-119, make no code changes; verify that updated_base and the copy bound remain within the validator-produced updated_key_bytes and updated_value_bytes after the widening fix.go/engine/hip/kernels/serving/kv_descriptor_append.hipinc-178-190 (1)
178-190: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDistinguish a legitimate page skip from a trim failure.
rocm_kv_descriptor_trim_pagereturnsfalsefor two different reasons:
- The page ends at or before
trim_start(line 48). The page holds no live tokens. Skipping it is correct.- The page straddles
trim_startbut the encoding is not row-interleaved, or the byte accounting does not match the stride model (lines 56-69). This is an error. The retained tokens cannot be addressed.Line 182 treats both cases as a skip. In case 2 the loop drops a page that still holds live tokens. The header written at lines 214-221 still reports
args.output_token_count. The descriptor then claims more tokens than its pages describe.The general fallback path at lines 375-391 already separates the two cases correctly. Line 378 skips the fully-trimmed page. Line 386 returns on a trim failure. Apply the same discrimination here.
Downstream,
rocm_valid_device_kv_attention_descriptoringo/engine/hip/kernels/serving/kv_device.hipinc(lines 50-73) validates each page in isolation. It does not verify that the pages cover the token range without a gap. A descriptor with a dropped page passes that validation, so attention reads a cache with missing tokens and produces a wrong result with no error signal.🐛 Proposed fix to separate the skip case from the failure case
for (uint32_t page_index = 0u; page_index < last_index; ++page_index) { const rocm_device_kv_page_descriptor *page = reinterpret_cast<const rocm_device_kv_page_descriptor *>(previous_base + ROCM_DEVICE_KV_DESCRIPTOR_HEADER_BYTES + page_index * ROCM_DEVICE_KV_DESCRIPTOR_PAGE_BYTES); + if (page->token_start + page->token_count <= args.trim_start) { + continue; + } rocm_device_kv_page_descriptor local_page {}; if (!rocm_kv_descriptor_trim_page(page, args.trim_start, &local_page)) { - continue; + return; } if (output_index >= output_last_index) { return; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/serving/kv_descriptor_append.hipinc` around lines 178 - 190, Update the page-processing loop around rocm_kv_descriptor_trim_page to distinguish a fully trimmed page from a trim failure: skip only pages ending at or before args.trim_start, but return immediately when a page straddles trim_start and trimming fails. Preserve the existing output bounds check and descriptor-page copying for successfully trimmed pages, matching the discrimination used in the general fallback path.go/engine/hip/kernels/serving/kv_device.hipinc-57-71 (1)
57-71: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winThe validator does not verify that the pages cover the token range.
The loop checks each page in isolation. Line 60 bounds each page within
args.token_count. No check compares the total covered tokens withargs.token_count, and no check verifies that pageN + 1starts where pageNends.Two malformed descriptors therefore pass validation:
- A descriptor with a gap. The header reports
token_counttokens, but the pages describe fewer. Attention then skips the missing tokens.- A descriptor with overlapping pages. Attention then attends to the same token twice.
Both produce a wrong result with no error signal. The append path in
go/engine/hip/kernels/serving/kv_descriptor_append.hipincat lines 178-190 can produce the gap case, so this is reachable.Accumulate the covered token count in the loop, and compare it with
args.token_countafter the loop.🛡️ Proposed fix to verify total coverage
const unsigned char *base = reinterpret_cast<const unsigned char *>(header); + uint64_t covered_tokens = 0; for (uint32_t page_index = 0; page_index < header->page_count; ++page_index) { const rocm_device_kv_page_descriptor *page = reinterpret_cast<const rocm_device_kv_page_descriptor *>(base + ROCM_DEVICE_KV_DESCRIPTOR_HEADER_BYTES + page_index * ROCM_DEVICE_KV_DESCRIPTOR_PAGE_BYTES); if (page->token_count == 0 || page->token_start + page->token_count > args.token_count ||return false; } + covered_tokens += page->token_count; } + if (covered_tokens != args.token_count) { + return false; + } return true;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/serving/kv_device.hipinc` around lines 57 - 71, Update the page-validation loop to accumulate each page’s covered token range and require contiguous, non-overlapping pages by verifying each page starts at the end of the previous page. After the loop, ensure the accumulated coverage equals args.token_count, while preserving the existing per-page validation in the descriptor validator.go/engine/hip/kernels/serving/kv_descriptor_append.hipinc-429-462 (1)
429-462: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winThe fused kernel can emit a descriptor that points at unencoded buffers.
This kernel runs three blocks. Blocks 0 and 1 encode the key and value tensors. Block 2 appends the page descriptor. Each block validates a different part of the packet:
- Blocks 0 and 1 validate
args.encodeat line 440. If that check fails, they return without writing any encoded data.- Block 2 validates
args.descriptorat line 458. If that check passes, it writes the page descriptor and the header.The two validations are independent. If
args.encodeis invalid andargs.descriptoris valid, block 2 publishes a descriptor whosekey_pointerandvalue_pointerreference buffers that were never written. The header reports the new token count. The host observes a successful launch.A later attention launch then reads uninitialised or stale KV data.
rocm_valid_device_kv_attention_descriptoringo/engine/hip/kernels/serving/kv_device.hipincchecks only descriptor metadata, so it cannot detect this.Validate both
args.encodeandargs.descriptorin every block before any block writes. Each block then makes the same accept or reject decision.🛡️ Proposed fix to make the accept decision uniform across blocks
if (args.version != ROCM_KV_ENCODE_TOKEN_VALUE_NORM_DESCRIPTOR_APPEND_LAUNCH_ARGS_VERSION || args.total_bytes != ROCM_KV_ENCODE_TOKEN_VALUE_NORM_DESCRIPTOR_APPEND_LAUNCH_ARGS_BYTES) { return; } + if (!rocm_valid_kv_encode_token_value_norm_args(args.encode) || + !rocm_valid_kv_descriptor_append_args(args.descriptor)) { + return; + } if (blockIdx.x < 2u) { - if (!rocm_valid_kv_encode_token_value_norm_args(args.encode)) { - return; - } const bool encode_key = blockIdx.x == 0;- if (!rocm_valid_kv_descriptor_append_args(args.descriptor)) { - return; - } rocm_kv_descriptor_append_execute(args.descriptor);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/serving/kv_descriptor_append.hipinc` around lines 429 - 462, Update rocm_kv_encode_token_value_norm_descriptor_append so every block validates both args.encode and args.descriptor before performing any encoding or descriptor append. Reuse the existing validation helpers, return immediately when either check fails, and remove the block-specific validation split so all three blocks make the same accept-or-reject decision.go/engine/hip/kernels/serving/kv_encode.hipinc-129-134 (1)
129-134: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winNon-power-of-two block launches can lose reduction results.
The strided reductions at the KV encode reduction sites fold the last active partial values in a loop with
stride = threads >> 1. WhenblockDim.xis not a power of two, values can remain in lanes that are never folded into lane 0, so the emitted scale/rms can be too small and quantisation clamps those KV values. Use a power-of-two block size, or fold the remainder tail into lane 0 before this reduction pattern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/serving/kv_encode.hipinc` around lines 129 - 134, Update the reduction logic in the KV encode kernel around the scratch-buffer loop so non-power-of-two block sizes cannot leave partial values out of lane 0. Either enforce a power-of-two block size for these launches or fold the remainder tail into lane 0 before the existing strided reduction, while preserving correct scale/RMS results for all supported block dimensions.go/engine/hip/kernels/serving/kv_encode.hipinc-436-452 (1)
436-452: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe global-scale Q4 path drops elements when
row_widthis odd, and it rescans the whole tensor per row.Two problems in this loop.
Correctness.
packed_countat line 436 covers the whole tensor. The loop at line 440 walks that whole range for every row, and line 442 keeps only the indices inside the current row.firstis always even, butrow_start = row * row_widthis odd for oddrow_widthon odd rows. The packed byte that straddles a row boundary is then handled by neither row correctly.Take
row_width == 3,row_count == 2,count == 6, sopacked_count == 3:
- Row 0 keeps
first0 and 2. Forfirst == 2,dim == 2anddim + 1 == 3is not less thanrow_width, so the high nibble ofpayload[1]is written as zero.- Row 1 keeps only
first == 4, givingdim == 1.- Element index 3, which is row 1 dimension 0, is never encoded. The high nibble of
payload[1]stays zero.The validator enforces an even width only for
ROCM_DEVICE_KV_DESCRIPTOR_ENCODING_Q4_ROWSandROCM_DEVICE_KV_DESCRIPTOR_ENCODING_Q4_ROWS_INTERLEAVEDat lines 60-63. PlainROCM_DEVICE_KV_DESCRIPTOR_ENCODING_Q4has no parity guard, and nothing forbidsrow_count > 1for it.The equivalent path in
rocm_encode_kv_token_tensorat lines 235-244 packs the flat stream and ignores row boundaries. That is correct for a single global scale, and it matches the size returned byrocm_device_kv_tensor_bytesat line 84 ofgo/engine/hip/kernels/serving/kv_device.hipinc.Cost. Each row scans all
packed_countindices and discards the ones outside the row. The loop performsrow_count * packed_countiterations wherepacked_countwould do.Drive the loop from a flat packed index, and derive the row from it.
🐛 Proposed fix to pack the flat stream and remove the per-row rescan
const uint32_t packed_count = (count + 1u) / 2u; for (uint32_t row = 0; row < row_count; ++row) { const uint32_t row_start = row * row_width; rocm_kv_encode_value_norm_prepare_row(input, row_start, head_dim, head_count, epsilon, scratch, head_rms); - for (uint32_t packed_index = tid; packed_index < packed_count; packed_index += threads) { - const uint32_t first = packed_index * 2u; - if (first < row_start || first >= row_start + row_width) { - continue; - } - const uint32_t dim = first - row_start; - const int low = rocm_quantize_kv_value(rocm_kv_encode_value_norm_value(input, row_start, dim, head_dim, head_rms), scale, encoding); - int high = 0; - if (dim + 1u < row_width) { - high = rocm_quantize_kv_value(rocm_kv_encode_value_norm_value(input, row_start, dim + 1u, head_dim, head_rms), scale, encoding); - } - payload[packed_index] = static_cast<unsigned char>(rocm_pack_signed_q4(low) | (rocm_pack_signed_q4(high) << 4)); - } + const uint32_t row_packed_begin = row_start / 2u; + const uint32_t row_packed_end = (row_start + row_width + 1u) / 2u; + for (uint32_t packed_index = row_packed_begin + tid; packed_index < row_packed_end; packed_index += threads) { + const uint32_t first = packed_index * 2u; + int low = 0; + int high = 0; + if (first >= row_start && first < row_start + row_width) { + low = rocm_quantize_kv_value(rocm_kv_encode_value_norm_value(input, row_start, first - row_start, head_dim, head_rms), scale, encoding); + } + if (first + 1u >= row_start && first + 1u < row_start + row_width && first + 1u < count) { + high = rocm_quantize_kv_value(rocm_kv_encode_value_norm_value(input, row_start, first + 1u - row_start, head_dim, head_rms), scale, encoding); + } + payload[packed_index] = static_cast<unsigned char>(rocm_pack_signed_q4(low) | (rocm_pack_signed_q4(high) << 4)); + } __syncthreads(); }A byte shared by two rows is still written twice under this fix. If plain Q4 with an odd
row_widthand multiple rows is not a supported shape, reject it inrocm_valid_kv_encode_token_value_norm_argsinstead:+ (args.value_encoding != ROCM_DEVICE_KV_DESCRIPTOR_ENCODING_Q4 || row_count == 1u || (value_width & 1u) == 0u) &&🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@go/engine/hip/kernels/serving/kv_encode.hipinc` around lines 436 - 452, Update the global-scale Q4 packing loop in rocm_encode_kv_token_value_norm to iterate each flat packed_index once, derive the corresponding row and dimensions from that index, and pack the flat tensor stream like rocm_encode_kv_token_tensor rather than rescanning packed_count for every row. If odd row_width with multiple rows cannot be safely represented because boundary bytes would be written twice, reject that shape in rocm_valid_kv_encode_token_value_norm_args.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: afafa3b6-c1af-4284-a04f-56ba9824bd47
📒 Files selected for processing (62)
go/engine/hip/hip_kernel_source_test.gogo/engine/hip/kernels/README.mdgo/engine/hip/kernels/activations/elementwise.hipincgo/engine/hip/kernels/activations/elementwise_validate.hipincgo/engine/hip/kernels/attention/attention_validate.hipincgo/engine/hip/kernels/attention/chunked.hipincgo/engine/hip/kernels/attention/device_kv.hipincgo/engine/hip/kernels/attention/heads.hipincgo/engine/hip/kernels/attention/rotary.hipincgo/engine/hip/kernels/attention/rotary_validate.hipincgo/engine/hip/kernels/attention/sdpa.hipincgo/engine/hip/kernels/attention/single_head.hipincgo/engine/hip/kernels/common/math.hipincgo/engine/hip/kernels/common/numeric.hipincgo/engine/hip/kernels/contract/launch_abi.hipincgo/engine/hip/kernels/embedding/lookup.hipincgo/engine/hip/kernels/embedding/lookup_validate.hipincgo/engine/hip/kernels/embedding/mean_pool.hipincgo/engine/hip/kernels/embedding/mean_pool_validate.hipincgo/engine/hip/kernels/embedding/rerank.hipincgo/engine/hip/kernels/embedding/rerank_validate.hipincgo/engine/hip/kernels/lora/projection.hipincgo/engine/hip/kernels/lora/projection_validate.hipincgo/engine/hip/kernels/matmul/projection.hipincgo/engine/hip/kernels/matmul/projection_validate.hipincgo/engine/hip/kernels/moe/routing.hipincgo/engine/hip/kernels/moe/routing_validate.hipincgo/engine/hip/kernels/norms/rms_norm.hipincgo/engine/hip/kernels/norms/rms_norm_fused_projection.hipincgo/engine/hip/kernels/norms/rms_norm_fused_projection_validate.hipincgo/engine/hip/kernels/norms/rms_norm_validate.hipincgo/engine/hip/kernels/quant/affine_dot.hipincgo/engine/hip/kernels/quant/affine_fused_validate.hipincgo/engine/hip/kernels/quant/affine_gelu_tanh.hipincgo/engine/hip/kernels/quant/affine_gelu_tanh_projection.hipincgo/engine/hip/kernels/quant/affine_projection.hipincgo/engine/hip/kernels/quant/affine_projection_fused.hipincgo/engine/hip/kernels/quant/affine_projection_tuned.hipincgo/engine/hip/kernels/quant/affine_projection_validate.hipincgo/engine/hip/kernels/quant/autoround.hipincgo/engine/hip/kernels/quant/autoround_validate.hipincgo/engine/hip/kernels/quant/gguf_projection.hipincgo/engine/hip/kernels/quant/gguf_validate.hipincgo/engine/hip/kernels/quant/jangtq_codebook.hipincgo/engine/hip/kernels/quant/jangtq_codebook_validate.hipincgo/engine/hip/kernels/rocm_kernels.hipgo/engine/hip/kernels/sampling/diffusion.hipincgo/engine/hip/kernels/sampling/diffusion_validate.hipincgo/engine/hip/kernels/sampling/greedy.hipincgo/engine/hip/kernels/sampling/greedy_validate.hipincgo/engine/hip/kernels/sampling/topk.hipincgo/engine/hip/kernels/sampling/topk_validate.hipincgo/engine/hip/kernels/serving/kv_descriptor_append.hipincgo/engine/hip/kernels/serving/kv_device.hipincgo/engine/hip/kernels/serving/kv_encode.hipincgo/engine/hip/kernels/serving/prefill_decode.hipincgo/engine/hip/kernels/serving/prefill_decode_validate.hipincgo/engine/hip/kernels/serving/tiny.hipincgo/engine/hip/kernels/serving/tiny_validate.hipincgo/engine/hip/kernels/serving/tiny_weights.hipincgo/engine/hip/kernels/training/losses.hipincgo/engine/hip/kernels/training/losses_validate.hipinc
| __device__ float rocm_block_reduce_sum(float value, float *scratch) | ||
| { | ||
| const uint32_t lane = threadIdx.x & 31u; | ||
| const uint32_t wave = threadIdx.x >> 5u; | ||
| const uint32_t wave_count = (blockDim.x + 31u) >> 5u; | ||
| for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { | ||
| value += rocm_shfl_down(value, stride, 32); | ||
| } | ||
| if (lane == 0u) { | ||
| scratch[wave] = value; | ||
| } | ||
| __syncthreads(); | ||
| float reduced = lane < wave_count ? scratch[lane] : 0.0f; | ||
| if (wave == 0u) { | ||
| for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { | ||
| reduced += rocm_shfl_down(reduced, stride, 32); | ||
| } | ||
| if (lane == 0u) { | ||
| scratch[0] = reduced; | ||
| } | ||
| } | ||
| __syncthreads(); | ||
| return scratch[0]; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Add a barrier after the final scratch[0] read to make repeated calls safe.
rocm_block_reduce_sum ends with a read of scratch[0] at line 25. The last barrier is at line 24, before that read. When the same scratch buffer is reused by a second call in the same kernel, the first write of the second call (scratch[wave] = value, line 12) has no barrier between it and the pending line 25 read of the first call. A fast thread can overwrite scratch[0] while a slow thread is still reading the previous total.
This affects rocm_moe_combine_norms (lines 160-161) and rocm_rms_norm_residual_add_norm (lines 193 and 208), which both call the reduction twice with the same partial array. The result is a non-deterministic wrong RMS value.
🐛 Proposed fix: read into a register, then synchronise
__syncthreads();
- return scratch[0];
+ const float result = scratch[0];
+ __syncthreads();
+ return result;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| __device__ float rocm_block_reduce_sum(float value, float *scratch) | |
| { | |
| const uint32_t lane = threadIdx.x & 31u; | |
| const uint32_t wave = threadIdx.x >> 5u; | |
| const uint32_t wave_count = (blockDim.x + 31u) >> 5u; | |
| for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { | |
| value += rocm_shfl_down(value, stride, 32); | |
| } | |
| if (lane == 0u) { | |
| scratch[wave] = value; | |
| } | |
| __syncthreads(); | |
| float reduced = lane < wave_count ? scratch[lane] : 0.0f; | |
| if (wave == 0u) { | |
| for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { | |
| reduced += rocm_shfl_down(reduced, stride, 32); | |
| } | |
| if (lane == 0u) { | |
| scratch[0] = reduced; | |
| } | |
| } | |
| __syncthreads(); | |
| return scratch[0]; | |
| } | |
| __device__ float rocm_block_reduce_sum(float value, float *scratch) | |
| { | |
| const uint32_t lane = threadIdx.x & 31u; | |
| const uint32_t wave = threadIdx.x >> 5u; | |
| const uint32_t wave_count = (blockDim.x + 31u) >> 5u; | |
| for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { | |
| value += rocm_shfl_down(value, stride, 32); | |
| } | |
| if (lane == 0u) { | |
| scratch[wave] = value; | |
| } | |
| __syncthreads(); | |
| float reduced = lane < wave_count ? scratch[lane] : 0.0f; | |
| if (wave == 0u) { | |
| for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { | |
| reduced += rocm_shfl_down(reduced, stride, 32); | |
| } | |
| if (lane == 0u) { | |
| scratch[0] = reduced; | |
| } | |
| } | |
| __syncthreads(); | |
| const float result = scratch[0]; | |
| __syncthreads(); | |
| return result; | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@go/engine/hip/kernels/norms/rms_norm.hipinc` around lines 3 - 26, Update
rocm_block_reduce_sum so it reads scratch[0] into a local register after the
existing final barrier, adds a __syncthreads() before returning that register,
and thereby prevents a subsequent call from overwriting the shared buffer while
another thread still reads it. Preserve the reduction and return value behavior
for callers such as rocm_moe_combine_norms and rocm_rms_norm_residual_add_norm.
| const uint32_t block_lane = threadIdx.x & (ROCM_Q8_1_BLOCK_SIZE - 1u); | ||
| const uint32_t workgroup_block = threadIdx.x / ROCM_Q8_1_BLOCK_SIZE; | ||
| const uint32_t blocks_per_workgroup = blockDim.x / ROCM_Q8_1_BLOCK_SIZE; | ||
| const uint32_t block_count = args.count / ROCM_Q8_1_BLOCK_SIZE; | ||
| for (uint32_t block_base = 0; block_base < block_count; block_base += blocks_per_workgroup) { |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Guard against blockDim.x < ROCM_Q8_1_BLOCK_SIZE to prevent an infinite loop.
Line 224 computes blocks_per_workgroup = blockDim.x / ROCM_Q8_1_BLOCK_SIZE with integer division. If the host launches this kernel with fewer than ROCM_Q8_1_BLOCK_SIZE threads per block, blocks_per_workgroup is 0. The loop at line 226 then never advances block_base, and block_count is at least 1 because the validator requires count > 0 and count % ROCM_Q8_1_BLOCK_SIZE == 0. The kernel hangs the device.
Other kernels in this layer already reject an unexpected block size. See rocm_projection at go/engine/hip/kernels/matmul/projection.hipinc line 15.
🛡️ Proposed fix
rocm_q8_1_block *q8_output = reinterpret_cast<rocm_q8_1_block *>(static_cast<uintptr_t>(args.q8_output_pointer));
+ if (blockDim.x < ROCM_Q8_1_BLOCK_SIZE) {
+ return;
+ }
const uint32_t block_lane = threadIdx.x & (ROCM_Q8_1_BLOCK_SIZE - 1u);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@go/engine/hip/kernels/norms/rms_norm.hipinc` around lines 222 - 226, Guard
the kernel entry path before the `block_base` loop so launches with `blockDim.x
< ROCM_Q8_1_BLOCK_SIZE` return immediately, preventing `blocks_per_workgroup`
from being zero. Keep the existing loop unchanged for valid block sizes, and
follow the early block-size rejection pattern used by `rocm_projection`.
| return args.version == ROCM_AUTOROUND_QUANTIZE_LAUNCH_ARGS_VERSION && | ||
| args.total_bytes == ROCM_AUTOROUND_QUANTIZE_LAUNCH_ARGS_BYTES && | ||
| args.weight_pointer != 0 && | ||
| args.packed_pointer != 0 && | ||
| args.scale_pointer != 0 && | ||
| args.rows > 0 && | ||
| args.cols > 0 && | ||
| args.group_size > 0 && | ||
| args.cols % args.group_size == 0 && | ||
| args.groups_per_row == args.cols / args.group_size && | ||
| rocm_valid_autoround_format(args.format_code, args.bits) && | ||
| args.weight_bytes == value_count * sizeof(float) && | ||
| args.packed_bytes == (packed_bits + 7u) / 8u && | ||
| args.scale_bytes == scale_count * sizeof(float) && | ||
| args.nsamples > 0 && | ||
| args.seqlen > 0 && | ||
| args.iters > 0; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
Groups that are not byte aligned cause a write race in rocm_autoround_quantize.
The shared root cause is a missing constraint in rocm_valid_autoround_quantize_args. The validator accepts any group_size that divides cols, and rocm_valid_autoround_format allows bits of 2, 4 or 8. It never requires (group_size * bits) % 8 == 0.
One thread owns one group. When group_size * bits is not a multiple of 8, two adjacent groups share a byte in packed. Two threads then touch the same byte with no atomics. Take bits == 4 with group_size == 3 and cols == 3: each group covers 12 bits, so group 1 starts mid-byte. Thread 1 clears the shared byte after thread 0 has already packed into it, and both threads perform a non-atomic read-modify-write OR. The packed output is corrupt and the result is not reproducible.
go/engine/hip/kernels/quant/autoround_validate.hipinc#L22-L38: add(args.group_size * args.bits) % 8 == 0to the validation chain. That makes every group start and end on a byte boundary, sopacked_group_bytescovers exactly the owned range.go/engine/hip/kernels/quant/autoround.hipinc#L38-L49: no change is needed once the constraint is added. The byte clear at Lines 38-40 and the OR inrocm_autoround_pack_signedthen stay inside the region owned by the thread.
🔒 Proposed fix in the validator
args.group_size > 0 &&
args.cols % args.group_size == 0 &&
+ (args.group_size * args.bits) % 8 == 0 &&
args.groups_per_row == args.cols / args.group_size &&📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return args.version == ROCM_AUTOROUND_QUANTIZE_LAUNCH_ARGS_VERSION && | |
| args.total_bytes == ROCM_AUTOROUND_QUANTIZE_LAUNCH_ARGS_BYTES && | |
| args.weight_pointer != 0 && | |
| args.packed_pointer != 0 && | |
| args.scale_pointer != 0 && | |
| args.rows > 0 && | |
| args.cols > 0 && | |
| args.group_size > 0 && | |
| args.cols % args.group_size == 0 && | |
| args.groups_per_row == args.cols / args.group_size && | |
| rocm_valid_autoround_format(args.format_code, args.bits) && | |
| args.weight_bytes == value_count * sizeof(float) && | |
| args.packed_bytes == (packed_bits + 7u) / 8u && | |
| args.scale_bytes == scale_count * sizeof(float) && | |
| args.nsamples > 0 && | |
| args.seqlen > 0 && | |
| args.iters > 0; | |
| return args.version == ROCM_AUTOROUND_QUANTIZE_LAUNCH_ARGS_VERSION && | |
| args.total_bytes == ROCM_AUTOROUND_QUANTIZE_LAUNCH_ARGS_BYTES && | |
| args.weight_pointer != 0 && | |
| args.packed_pointer != 0 && | |
| args.scale_pointer != 0 && | |
| args.rows > 0 && | |
| args.cols > 0 && | |
| args.group_size > 0 && | |
| args.cols % args.group_size == 0 && | |
| (args.group_size * args.bits) % 8 == 0 && | |
| args.groups_per_row == args.cols / args.group_size && | |
| rocm_valid_autoround_format(args.format_code, args.bits) && | |
| args.weight_bytes == value_count * sizeof(float) && | |
| args.packed_bytes == (packed_bits + 7u) / 8u && | |
| args.scale_bytes == scale_count * sizeof(float) && | |
| args.nsamples > 0 && | |
| args.seqlen > 0 && | |
| args.iters > 0; |
📍 Affects 2 files
go/engine/hip/kernels/quant/autoround_validate.hipinc#L22-L38(this comment)go/engine/hip/kernels/quant/autoround.hipinc#L38-L49
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@go/engine/hip/kernels/quant/autoround_validate.hipinc` around lines 22 - 38,
Update rocm_valid_autoround_quantize_args in
go/engine/hip/kernels/quant/autoround_validate.hipinc:22-38 to require
(args.group_size * args.bits) % 8 == 0 alongside the existing group-size checks.
This ensures each packed group occupies whole bytes; rocm_autoround_quantize and
rocm_autoround_pack_signed in go/engine/hip/kernels/quant/autoround.hipinc:38-49
require no direct changes because the validator constraint prevents overlapping
byte access.
#23) * refactor(engine/metal): the kernel pile takes the family shape hip set engine/hip's kernels/ was restructured into the QuixiCore family taxonomy (PR #22); engine/metal was a flat pile of 47 .metal files + 2 shared _impl.h in one directory. Same treatment, same vocabulary — the two engines now read alike: activations/ 5 attention/ 14 embedding/ 4 megakernel/ 4 moe/ 6+2h norms/ 3 probes/ 2 quant/ 7 sampling/ 1 training/ 1 experimental/ 1 Unlike the HIP amalgamation TU, this is a SORT, not a split: every file moved byte-identically (git records all 49 as R100; per-file sha256 old->new all match), no shader edited, no file or kernel function renamed. Classification is by what the kernel does, not by its filename prefix — lthn_moe_combine_norms is the MoE block's norm/combine tail and sits in norms/ exactly as rocm_moe_combine does on engine/hip; lthn_gather_qmv + lthn_gelu_qmv are the affine expert-route projections and sit in moe/ (which also keeps each shared _impl.h beside its consumers, since the quoted #include resolves relative to the includer). Two families have no HIP counterpart and are named for what they are: megakernel/ (the grid-barrier whole-stage dispatches, an Apple-Silicon lever) and probes/ (grid-sync + cross-threadgroup coherency, compiled into the library because the tests dispatch them). contract/ and lora/ do not exist here; common/ and matmul/ are absent by observation — there is no engine-wide numeric header, and every projection kernel here is quantised. Build glue: task metallib:kernels now collects sources RECURSIVELY over kernels/, excluding experimental/ (which was never in the shipped library and still is not) — .metal files are independent compilation units linked into one metallib, so the tree carries no ordering meaning and a new family directory needs no build change. A duplicate-basename guard fails the build loudly: the .air objects are named by basename, so two families holding one filename would silently drop a kernel. Receipts, this Mac (M-series, MetalToolchain v17.6, macOS 26): - 47 sources before, 47 after; no dupes. - lthn_kernels.metallib is REPRODUCIBLE: the baseline rebuild reproduced the 2026-07-22 artefact byte for byte (sha256 6be8385f0b63f79d…762207d), so the metallib hash is a usable no-regression proof rather than a moving target. - 209 exported functions via `xcrun metal-nm` — the parity instrument, and the reason kernels/README.md carries no hand-maintained symbol list. Co-Authored-By: Virgil <virgil@lethean.io> * docs(engine/metal): point the Go comment references at the family they moved to 26 doc comments across 17 files name a kernel source by path (kernels/lthn_tq_kv.metal, "compiled from kernels/*.metal", …). Nothing on the Go side READS those paths — the engine resolves kernels by function name off the loaded library (customLibrary.NewFunctionWithName), which is why the move needed no launcher change — but "path is documentation" only holds while the path is true. Comment text only; substitution was driven by the exact old->new table the move itself produced, and every changed line was read back in the diff. go build + go vet green. Co-Authored-By: Virgil <virgil@lethean.io>
…d on Windows) (#30) v0.12.1 carries the two core/go fixes this repo's windows lane has been blocked on, and which no amount of work here could reach: - #21 — io.Local was local.New("/"), and "/" is not an absolute root on Windows, it names the CURRENT DRIVE. An absolute "C:\Users\..." under a D:-rooted medium gave "path escapes from parent"; a relative "testdata/x.json" resolved against the drive root instead of the cwd and gave "The system cannot find the path specified". 30 of the 34 remaining packages report one shape or the other. - #22 — App.Find carried the same three defects go-process v0.16.2 fixed: no %PATHEXT% expansion, a mode&0111 test Windows can never satisfy (os.Stat synthesises 0666/0444, so it rejected every file including git.exe), and a path-vs-name check that missed '/'. 3 packages report `Program.Find: "git": not found in PATH`. Receipts — macOS: GOWORK=off go test -count=1 ./... exit=0, 176 packages ok GOWORK=off go build ./... · go vet ./... clean cli: go test -count=1 ./tui/... ok 6.510s examples: go build ./... · go vet ./... clean Only go/ is bumped. cli/ pins v0.12.0 and examples/ v0.11.0, but both build through the workspace, which already selects the highest version across modules — the skew predates this and is not widened by it. Lane state going in: 34 (main, run 31251799117). The windows number is this lane's own and is posted from the run, not predicted here.
1 participant
The 17,090-line
rocm_kernels.hipbecomes a 72-line amalgamation TU over 59 contiguous fragments inkernels/<family>/(contract · common · quant · attention · serving · sampling · norms · moe · embedding · matmul · activations · training · lora), following the QuixiCore taxonomy adapted to this engine (thevariants/<target>axis is deliberately unused — one source compiles AMD/NVIDIA/CPU; the README records where it goes if arch-specific kernels arrive).Byte-identity receipts: assembled source sha256 identical to HEAD's mega file (
454c04d0…); all 142 exported kernel bodies identical; the only diff is two manifest banner comments. A manifest-coverage test fails on any fragment on disk but unwired, wired twice, or missing its licence header (negative control run included).Validation: Mac
GOWORK=off go build ./...+ vet clean; homelab compiles at identical output sizes on ROCm gfx1101, CUDA 12.8, and HIP-CPU x86_64;go test ./engine/hipok; the 42TestHIPKernelSource_tests pass over the new layout.task test's engine/metal failure is pre-existing at the base commit (symlinked config.json rejected by the sandboxed reader) and untouched by this diff.🤖 Generated with Claude Code
Co-Authored-By: Virgil virgil@lethean.io
Summary by CodeRabbit
New Features
Documentation
Tests