feat(cpu): add Qwen3.5 0.8B mobile CPU support - #690
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds CPU Qwen3.5 hybrid inference with GDN and full-attention layers, tokenizer and CLI support, KAI quantization and conversion validation, generation metrics, mobile converter validation, contiguous-copy updates, and tests. ChangesQwen3.5 runtime and example
Conversion and model-file validation
Generation performance metrics
Contiguous tensor operation
Mobile module loading
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant Qwen3_5Tokenizer
participant Qwen3_5ForCausalLM
participant Qwen3_5Text
participant CPUKernels
CLI->>Qwen3_5Tokenizer: convert prompt to token ids
CLI->>Qwen3_5ForCausalLM: reset state and generate
Qwen3_5ForCausalLM->>Qwen3_5Text: run hybrid decoder with KV cache
Qwen3_5Text->>CPUKernels: execute GDN convolution and gated delta rule
CPUKernels-->>Qwen3_5Text: return recurrent-layer outputs
Qwen3_5Text-->>Qwen3_5ForCausalLM: return logits
Qwen3_5ForCausalLM-->>CLI: stream generated tokens
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
|
The Android, Linux x86_64, and macOS workflows for this fork PR are currently |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mllm/models/ARGeneration.cpp (1)
152-160: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
prefillEventStartTimePoint()zeroesar_prefill_tokens_right after it's set — batchgenerate()always reports 0 prefill tokens.
ar_prefill_tokens_is assigned frompast["sequence"].shape()[1]at line 155, thenprefillEventStartTimePoint()is called at line 156, which unconditionally resetsar_prefill_tokens_ = 0(see the rewritten helper, lines 534-546). The value set on line 155 is immediately discarded, soperfStats().prefill_tokens— and thereforeprefill_tokens_per_secinperfSummary()— will always read 0 for the batchgenerate()API.Compare with the correct ordering used elsewhere in this same diff:
ARGenerationChatIterator::step()callsprefillEventStartTimePoint()first, then setsar_prefill_tokens_afterforward()returns (lines 81, 90).streamGenerate()does the same (lines 230, 239).Only
generate()has the assignment before the reset. This isn't caught by the new tests, sinceARGenerationTest.cpponly exercises thechat()iterator path.🐛 Proposed fix
if (i == 0) { - if (past.count("sequence") > 0) { ar_prefill_tokens_ = past["sequence"].shape()[1]; } prefillEventStartTimePoint(); + if (past.count("sequence") > 0) { ar_prefill_tokens_ = past["sequence"].shape()[1]; } } else { decodeEventStartTimePoint(); }🤖 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 `@mllm/models/ARGeneration.cpp` around lines 152 - 160, Reorder the first-iteration logic in generate() so prefillEventStartTimePoint() runs before assigning ar_prefill_tokens_ from past["sequence"].shape()[1]. Preserve the existing decode timing path and ensure the batch generate() API retains the measured prefill token count for perfStats() and perfSummary().
🧹 Nitpick comments (10)
pymllm/mobile/__init__.py (1)
13-19: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the module-level
__getattr__contract.The docstring only describes the purpose; document the
nameparameter, returned module, andAttributeErrorraised for unsupported attributes.As per coding guidelines, public APIs, classes, and functions must have clear docstrings or comments explaining purpose, parameters, returns, and errors.
Proposed docstring
def __getattr__(name: str): - """Load optional subsystems only when a caller requests them.""" + """Lazily load optional subsystems. + + Args: + name: Requested module attribute. + Returns: + The requested optional subsystem module. + Raises: + AttributeError: If the attribute is not supported. + """🤖 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 `@pymllm/mobile/__init__.py` around lines 13 - 19, Expand the module-level __getattr__ docstring to document the name parameter, the lazily returned service module, and the AttributeError raised for unsupported attribute names. Keep the existing lazy import and caching behavior unchanged.Source: Coding guidelines
pymllm/mobile/utils/mllm_convertor.py (1)
58-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRedundant re-filtering:
load_modelalready appliesinclude_prefixes.
convertor.load_model(args.input_path, args.include_prefix)already filters tensors via its internalshould_load()predicate using the exact same prefix-matching semantics. Re-filteringparamshere duplicates that logic; if the two implementations ever diverge, results could silently disagree.♻️ Proposed simplification
# Get params params = convertor.load_model(args.input_path, args.include_prefix) if args.include_prefix: - params = { - name: tensor - for name, tensor in params.items() - if any(name.startswith(prefix) for prefix in args.include_prefix) - } if not params: parser.error( "--include_prefix did not match any parameters; refusing to write an empty model" ) if args.verbose: print( f"Kept {len(params)} parameters matching prefixes: " + ", ".join(args.include_prefix) )🤖 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 `@pymllm/mobile/utils/mllm_convertor.py` around lines 58 - 74, Remove the redundant parameter filtering comprehension after convertor.load_model in the conversion flow. Keep the empty-result validation and verbose reporting based on the params returned by load_model, which already applies args.include_prefix with the canonical matching logic.mllm/ffi/Extension.cc (1)
383-383: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueMove the mid-file include into the top-of-file include block.
Extension.cckeeps the normal include block at the top, but this header is still included inline before the registration block at line 383. Keep headers grouped at the top of the file.🤖 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 `@mllm/ffi/Extension.cc` at line 383, Move the kai_w4a32_pack.hpp include from its current mid-file location into Extension.cc’s existing top-of-file include block, leaving the registration block and surrounding implementation unchanged.Source: Coding guidelines
mllm/models/ARGeneration.cpp (1)
64-134: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftTiming bookkeeping is duplicated near-identically across
step(),generate(), andstreamGenerate().All three implement the same prefill/decode/first-token/EOS timing sequence with only minor per-API differences (device handling, callback vs. return). This duplication is exactly why the ordering bug in
generate()(see the comment on lines 154-159) could slip in — the three copies drifted out of sync. Consider extracting a shared private helper (e.g.,recordStepTiming(bool is_first_step, const ARGenerationOutputPast& pre_forward_input)) that all three call sites use, to guarantee a single source of truth for event ordering.Also applies to: 138-211, 213-279
🤖 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 `@mllm/models/ARGeneration.cpp` around lines 64 - 134, The timing bookkeeping is duplicated across ARGenerationChatIterator::step(), generate(), and streamGenerate(), allowing event ordering to diverge. Extract a shared private helper for prefill, decode, first-token, and EOS timing, passing the step context needed by each API, then replace each duplicated timing sequence with calls to that helper while preserving their existing device, callback, and return behavior.tests/core/ARGenerationTest.cpp (1)
1-133: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo test coverage for
ARGeneration::generate()orstreamGenerate()performance stats.All new tests exercise only the
chat()iterator path. This is why thear_prefill_tokens_reset-ordering bug ingenerate()(flagged inmllm/models/ARGeneration.cpplines 154-159) wasn't caught. Consider adding at least onegenerate()-based test assertingstats.prefill_tokensmatches the input sequence length.🤖 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 `@tests/core/ARGenerationTest.cpp` around lines 1 - 133, Add a performance-statistics test using ARGeneration::generate(), not only runChat()/chat(), and assert the resulting perfStats().prefill_tokens equals the input sequence length. Reuse FakeGeneration and makeInput, while preserving the existing assertions and test coverage.mllm/models/ARGeneration.hpp (1)
27-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new public stats API and mark
perfStats()[[nodiscard]].
ARGenerationPerformanceStatsandperfStats()are new public surface with no docstrings explaining field units/semantics (e.g., whatvalid/completedgate,_ussuffix meaning), and static analysis flags the accessor as missing[[nodiscard]].As per coding guidelines, "Ensure public APIs, classes, and functions have clear docstrings or comments explaining purpose, parameters, returns, and errors."
📝 Proposed fix
+// Snapshot of generation timing/counters produced by ARGeneration::perfStats(). +// All *_us fields are microseconds; `valid` is false until prefill, prefill-end, +// and first-token events have all been recorded for the current request. struct ARGenerationPerformanceStats { bool valid = false; bool completed = false; @@ - ARGenerationPerformanceStats perfStats() const; + [[nodiscard]] ARGenerationPerformanceStats perfStats() const;Also applies to: 109-110
🤖 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 `@mllm/models/ARGeneration.hpp` around lines 27 - 38, Document the public ARGenerationPerformanceStats struct and each field’s semantics, including that valid/completed gate the measurements, _us fields are microseconds, and token/step counters’ meanings. Add clear documentation for the public perfStats() accessor covering its return value, then mark perfStats() [[nodiscard]] so callers cannot ignore the result.Sources: Coding guidelines, Linters/SAST tools
mllm/backends/cpu/kernels/common/kai_w4a32_pack.hpp (1)
9-25: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd docstrings for the new shared packing API.
KaiW4A32Tile,kaiW4A32TileFromName,kaiW4A32PackedSize, andkaiW4A32QuantizeAndPackare now a cross-module contract (consumed frommllm/ffi/Extension.ccandkai.cpp), but only a single top-of-file comment exists. Document parameter ordering (out_channels/in_channels), the meaning oftile_name, and thestd::invalid_argumenterror contract for unsupported/invalid inputs.As per coding guidelines, "Ensure public APIs, classes, and functions have clear docstrings or comments explaining purpose, parameters, returns, and errors."
🤖 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 `@mllm/backends/cpu/kernels/common/kai_w4a32_pack.hpp` around lines 9 - 25, Add clear API documentation for KaiW4A32Tile and the kaiW4A32TileFromName, kaiW4A32PackedSize, and kaiW4A32QuantizeAndPack declarations. Describe tile_name, preserve out_channels/in_channels ordering, document return or output behavior and relevant parameters, and state that unsupported or invalid inputs throw std::invalid_argument.Source: Coding guidelines
examples/qwen3_5/quant_cfg_0.8B_w4a32_kai.json (1)
80-105: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider leaving the GDN gate projections in float32.
in_proj_a/in_proj_bare[16, 1024]— negligible memory savings, but their outputs drive the decay/beta gates in the recurrent update, where 4-bit weight error accumulates across the sequence. Keeping them unquantized costs ~128 KB total and removes a plausible quality risk.🤖 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 `@examples/qwen3_5/quant_cfg_0.8B_w4a32_kai.json` around lines 80 - 105, Remove the quantization entries for the in_proj_a and in_proj_b patterns from the configuration so these GDN gate projections remain in float32. Leave the surrounding quantization rules unchanged.mllm/models/qwen3_5/tokenization_qwen3_5.hpp (1)
255-257: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the placeholder length instead of hardcoding
12.- size_t pos = applied_string.find("{{{prompt}}}"); + static constexpr std::string_view kPromptPlaceholder = "{{{prompt}}}"; + size_t pos = applied_string.find(kPromptPlaceholder); if (pos == std::string::npos) { throw std::runtime_error("Qwen3.5 message template is missing the prompt placeholder"); } - applied_string.replace(pos, 12, message.prompt); + applied_string.replace(pos, kPromptPlaceholder.size(), message.prompt);🤖 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 `@mllm/models/qwen3_5/tokenization_qwen3_5.hpp` around lines 255 - 257, Update the placeholder replacement in the message-template logic around applied_string and "{{{prompt}}}" to derive the replacement length from the placeholder string itself rather than hardcoding 12, while preserving the existing missing-placeholder validation and replacement behavior.examples/qwen3_5/validate_checkpoint.py (1)
18-18: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnderscore-prefixed helper is imported cross-module.
validate_converted_model.pyline 16 imports_expected_text_shapesfrom this module, which contradicts the private naming. Renaming it toexpected_text_shapes(and updating the import) makes the shared contract explicit.🤖 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 `@examples/qwen3_5/validate_checkpoint.py` at line 18, Rename the helper _expected_text_shapes to expected_text_shapes and update the corresponding import and references in validate_converted_model.py, preserving its existing behavior and signature.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@examples/qwen3_5/main.cpp`:
- Around line 92-95: Update the output loop over model.chat results to stop
using std::wcout for detokenized text. Convert
tokenizer.detokenize(step.cur_token_id) with preprocessor::wideString2Utf8String
and emit the resulting UTF-8 text via fmt::print, preserving token ID
diagnostics and flush behavior while keeping stdout narrow-oriented.
In `@mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp`:
- Around line 78-97: Add an ARM-only test that iterates over every tile name
supported by kaiW4A32TileFromName and compares its returned nr, kr, and sr
values with the corresponding ARM KAI ukernel get_nr, get_kr, and get_sr values.
Reuse the existing tile-to-ukernel mapping and test infrastructure, and retain
the current non-ARM determinism test unchanged.
- Around line 26-30: Update truncateToBFloat16 to convert the float bit pattern
to bfloat16 using round-to-nearest-even rather than simply shifting away the
lower 16 bits. Preserve the returned uint16_t representation and ensure the
conversion handles carry into the upper half correctly before the per-block
scale is passed with kai_dt_bf16.
In `@mllm/models/qwen3_5/configuration_qwen3_5.hpp`:
- Around line 127-158: Add [[nodiscard]] to the const accessor methods
rotary_dim(), isFullAttentionLayer(), numFullAttentionLayers(), and
numGDNLayers() in the configuration class, preserving their existing signatures
and behavior.
- Around line 74-78: Validate explicitly provided linear_impl_type values in the
configuration-loading logic before assigning the result from
aops::str2LinearImplTypes. If parsing an explicitly supplied value yields
LinearImplTypes::kDefault, reject the configuration with the existing
configuration-error mechanism instead of silently accepting it; preserve the
default behavior when the field is absent.
In `@mllm/models/qwen3_5/modeling_qwen3_5.hpp`:
- Line 561: Update the Qwen3_5 generation stopping logic around eos_token_id_ so
the configured Qwen3_5Config::eos_token_id is honored in addition to
im_end_token_id, allowing either token to terminate generation. Preserve the
existing im_end_token_id behavior while ensuring the parsed configuration value
is no longer unused.
- Line 249: Initialize the layer_idx_ member in Qwen3_5FullAttention with a safe
default value so default-constructed instances cannot pass an indeterminate
index to updateKVCache. Preserve Qwen3_5Text’s later assignment for explicitly
configured layers.
In `@mllm/models/qwen3_5/tokenization_qwen3_5.hpp`:
- Around line 144-152: Update the unmatched branch in the tokenization loop
around qwen3_5TokenizerMatchPattern so it emits the current single character as
its own piece before advancing pos. Preserve the existing matched-pattern
behavior and ensure every unmatched character remains in the token stream for
lossless byte-wise encoding.
In `@pymllm/mobile/convertor/__init__.py`:
- Around line 117-127: Update the torch checkpoint loading in load_model() to
call torch.load with weights_only=True, preserving map_location="cpu" and the
existing prefix filtering and return behavior.
In `@pymllm/mobile/quantize/solver.py`:
- Around line 92-104: Make the collision validation around planned_names
order-independent by separating name removal from output-name collision checks.
First process every replace=True payload in param_groups to remove its input
names from the planned set, then perform the existing output collision checks
and additions in a second pass, preserving the ValueError behavior for genuine
collisions.
In `@tests/cpu/ContiguousOpTest.cpp`:
- Around line 13-38: Add a rank-0 scalar tensor case to CPUContiguousOpTest,
exercising Tensor::contiguous() and validating the scalar shape and copied value
so the ndim == 0 fast path in CPUContiguousOp::forward is covered. Keep the
existing contiguous and strided-view assertions unchanged.
In `@tests/cpu/Qwen35GDNTest.cpp`:
- Around line 199-255: Restore the process-wide CPU thread count after
ParallelBatchValueHeadsMatchSerialBitwise completes. Capture the existing value
before setCpuOpThreads(kThreadCount), then restore it after the comparisons,
ensuring restoration also occurs if an assertion aborts the test path.
---
Outside diff comments:
In `@mllm/models/ARGeneration.cpp`:
- Around line 152-160: Reorder the first-iteration logic in generate() so
prefillEventStartTimePoint() runs before assigning ar_prefill_tokens_ from
past["sequence"].shape()[1]. Preserve the existing decode timing path and ensure
the batch generate() API retains the measured prefill token count for
perfStats() and perfSummary().
---
Nitpick comments:
In `@examples/qwen3_5/quant_cfg_0.8B_w4a32_kai.json`:
- Around line 80-105: Remove the quantization entries for the in_proj_a and
in_proj_b patterns from the configuration so these GDN gate projections remain
in float32. Leave the surrounding quantization rules unchanged.
In `@examples/qwen3_5/validate_checkpoint.py`:
- Line 18: Rename the helper _expected_text_shapes to expected_text_shapes and
update the corresponding import and references in validate_converted_model.py,
preserving its existing behavior and signature.
In `@mllm/backends/cpu/kernels/common/kai_w4a32_pack.hpp`:
- Around line 9-25: Add clear API documentation for KaiW4A32Tile and the
kaiW4A32TileFromName, kaiW4A32PackedSize, and kaiW4A32QuantizeAndPack
declarations. Describe tile_name, preserve out_channels/in_channels ordering,
document return or output behavior and relevant parameters, and state that
unsupported or invalid inputs throw std::invalid_argument.
In `@mllm/ffi/Extension.cc`:
- Line 383: Move the kai_w4a32_pack.hpp include from its current mid-file
location into Extension.cc’s existing top-of-file include block, leaving the
registration block and surrounding implementation unchanged.
In `@mllm/models/ARGeneration.cpp`:
- Around line 64-134: The timing bookkeeping is duplicated across
ARGenerationChatIterator::step(), generate(), and streamGenerate(), allowing
event ordering to diverge. Extract a shared private helper for prefill, decode,
first-token, and EOS timing, passing the step context needed by each API, then
replace each duplicated timing sequence with calls to that helper while
preserving their existing device, callback, and return behavior.
In `@mllm/models/ARGeneration.hpp`:
- Around line 27-38: Document the public ARGenerationPerformanceStats struct and
each field’s semantics, including that valid/completed gate the measurements,
_us fields are microseconds, and token/step counters’ meanings. Add clear
documentation for the public perfStats() accessor covering its return value,
then mark perfStats() [[nodiscard]] so callers cannot ignore the result.
In `@mllm/models/qwen3_5/tokenization_qwen3_5.hpp`:
- Around line 255-257: Update the placeholder replacement in the
message-template logic around applied_string and "{{{prompt}}}" to derive the
replacement length from the placeholder string itself rather than hardcoding 12,
while preserving the existing missing-placeholder validation and replacement
behavior.
In `@pymllm/mobile/__init__.py`:
- Around line 13-19: Expand the module-level __getattr__ docstring to document
the name parameter, the lazily returned service module, and the AttributeError
raised for unsupported attribute names. Keep the existing lazy import and
caching behavior unchanged.
In `@pymllm/mobile/utils/mllm_convertor.py`:
- Around line 58-74: Remove the redundant parameter filtering comprehension
after convertor.load_model in the conversion flow. Keep the empty-result
validation and verbose reporting based on the params returned by load_model,
which already applies args.include_prefix with the canonical matching logic.
In `@tests/core/ARGenerationTest.cpp`:
- Around line 1-133: Add a performance-statistics test using
ARGeneration::generate(), not only runChat()/chat(), and assert the resulting
perfStats().prefill_tokens equals the input sequence length. Reuse
FakeGeneration and makeInput, while preserving the existing assertions and test
coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: b70de051-ffae-44bb-9220-17089c52f9c2
📒 Files selected for processing (34)
examples/CMakeLists.txtexamples/qwen3_5/CMakeLists.txtexamples/qwen3_5/README.mdexamples/qwen3_5/config_0.8B_w4a32_kai.jsonexamples/qwen3_5/main.cppexamples/qwen3_5/quant_cfg_0.8B_w4a32_kai.jsonexamples/qwen3_5/validate_checkpoint.pyexamples/qwen3_5/validate_converted_model.pymllm/backends/cpu/CMakeLists.txtmllm/backends/cpu/kernels/arm/linear/kai.cppmllm/backends/cpu/kernels/common/gdn/gated_delta_net.cppmllm/backends/cpu/kernels/common/gdn/gated_delta_net.hppmllm/backends/cpu/kernels/common/kai_w4a32_pack.cppmllm/backends/cpu/kernels/common/kai_w4a32_pack.hppmllm/backends/cpu/ops/ContiguousOp.cppmllm/ffi/Extension.ccmllm/models/ARGeneration.cppmllm/models/ARGeneration.hppmllm/models/qwen3_5/configuration_qwen3_5.hppmllm/models/qwen3_5/modeling_qwen3_5.hppmllm/models/qwen3_5/tokenization_qwen3_5.hpppymllm/mobile/__init__.pypymllm/mobile/convertor/__init__.pypymllm/mobile/convertor/model_file_v2.pypymllm/mobile/quantize/kai/w4a32.pypymllm/mobile/quantize/solver.pypymllm/mobile/utils/mllm_convertor.pytests/core/ARGenerationTest.cpptests/core/CMakeLists.txttests/cpu/CMakeLists.txttests/cpu/ContiguousOpTest.cpptests/cpu/KaiW4A32PackTest.cpptests/cpu/Qwen35GDNTest.cpptests/cpu/Qwen35TokenizerTest.cpp
| uint16_t truncateToBFloat16(float value) { | ||
| uint32_t bits = 0; | ||
| std::memcpy(&bits, &value, sizeof(bits)); | ||
| return static_cast<uint16_t>(bits >> 16); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file location =="
git ls-files | rg 'mllm/backends/cpu/kernels/common/kai_w4a32_pack\.cpp|kai_w4a32_pack' || true
echo "== inspect target file =="
if [ -f mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp ]; then
nl -ba mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp | sed -n '1,120p'
fi
echo "== search related bf16 helpers/usages =="
rg -n "truncateToBFloat16|roundToBFloat16|cast_bf16_f32|kai_cast_bf16_f32|BFloat16|bf16|float_to_bf16" mllm pymllm mllm-cli 2>/dev/null || true
echo "== deterministic f32->bf16 compare for small range =="
python3 - <<'PY'
import struct
def trunc_bf16(value):
bits = struct.unpack('<I', struct.pack('<f', value))[0]
return bits >> 16
def rnd_bf16(value):
bits = struct.unpack('<I', struct.pack('<f', value))[0]
rounding_bias = ((bits >> 16) & 1U) + 0x7FFFU
out = (bits + rounding_bias) >> 16
# Handle inf/nan preserved through arithmetic with bias? keep same as trunc?
return out
def uf16_to_float(uf16):
return struct.unpack('<f', struct.pack('<I', uf16 << 16))[0]
values = list(struct.unpack("<2147483648f", bytes([0x7f]*4)))[::2147483647//2] if False else []
samples = [
1.5, 2.5, 3.5, 4.5, 5.5,
-1.5, -2.5, -3.5, -4.5, -5.5,
0.5, 1.0, 7.0, 8.0, 16.0,
0.000977517, 0.000977518,
]
for v in samples:
b_trunc = trunc_bf16(v)
b_rnd = rnd_bf16(v)
print(v, u"{:016b}".format(b_trunc), uf16_to_float(b_trunc), u"{:016b}".format(b_rnd), uf16_to_float(b_rnd))
PYRepository: UbiquitousLearning/mllm
Length of output: 351
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file location =="
git ls-files | rg 'mllm/backends/cpu/kernels/common/kai_w4a32_pack\.cpp|kai_w4a32_pack' || true
echo "== inspect target file =="
if [ -f mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp ]; then
nl -ba mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp | sed -n '1,120p'
fi
echo "== search related bf16 helpers/usages =="
rg -n "truncateToBFloat16|roundToBFloat16|cast_bf16_f32|kai_cast_bf16_f32|BFloat16|bf16|float_to_bf16" mllm pymllm mllm-cli 2>/dev/null || true
echo "== deterministic f32->bf16 compare for selected cases =="
python3 - <<'PY'
import struct
def trunc_bf16(value):
bits = struct.unpack('<I', struct.pack('<f', value))[0]
return bits >> 16
def rnd_bf16(value):
bits = struct.unpack('<I', struct.pack('<f', value))[0]
rounding_bias = ((bits >> 16) & 1U) + 0x7FFFU
return (bits + rounding_bias) >> 16
def uf16_to_float(uf16):
return struct.unpack('<f', struct.pack('<I', uf16 << 16))[0]
samples = [
1.5, 2.5, 3.5, 4.5, 5.5,
-1.5, -2.5, -3.5, -4.5, -5.5,
0.5, 1.0, 7.0, 8.0, 16.0,
0.000977517, 0.000977518,
]
for v in samples:
b_trunc = trunc_bf16(v)
b_rnd = rnd_bf16(v)
print(v, f"{b_trunc:016b}", uf16_to_float(b_trunc), f"{b_rnd:016b}", uf16_to_float(b_rnd))
PYRepository: UbiquitousLearning/mllm
Length of output: 351
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== inspect target file =="
if [ -f mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp ]; then
awk '{ printf "%6d\t%s\n", NR, $0 }' mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp | sed -n '1,140p'
fi
echo "== inspect header (if relevant) =="
if [ -f mllm/backends/cpu/kernels/common/kai_w4a32_pack.hpp ]; then
awk '{ printf "%6d\t%s\n", NR, $0 }' mllm/backends/cpu/kernels/common/kai_w4a32_pack.hpp | sed -n '1,140p'
fi
echo "== search related bf16 helpers/usages =="
rg -n "truncateToBFloat16|roundToBFloat16|cast_bf16_f32|kai_cast_bf16_f32|BFloat16|bf16|float_to_bf16" mllm pymllm mllm-cli 2>/dev/null || true
echo "== deterministic f32->bf16 compare for selected cases =="
python3 - <<'PY'
import struct
def trunc_bf16(value):
bits = struct.unpack('<I', struct.pack('<f', value))[0]
return bits >> 16
def rnd_bf16(value):
bits = struct.unpack('<I', struct.pack('<f', value))[0]
rounding_bias = ((bits >> 16) & 1U) + 0x7FFFU
return (bits + rounding_bias) >> 16
def uf16_to_float(uf16):
return struct.unpack('<f', struct.pack('<I', uf16 << 16))[0]
samples = [
1.5, 2.5, 3.5, 4.5, 5.5,
-1.5, -2.5, -3.5, -4.5, -5.5,
0.5, 1.0, 7.0, 8.0, 16.0,
0.000977517, 0.000977518,
]
for v in samples:
b_trunc = trunc_bf16(v)
b_rnd = rnd_bf16(v)
print(v, f"{b_trunc:016b}", uf16_to_float(b_trunc), f"{b_rnd:016b}", uf16_to_float(b_rnd))
PYRepository: UbiquitousLearning/mllm
Length of output: 11009
🌐 Web query:
KleidiAI kai_cast_bf16_f32 implementation bf16 conversion
💡 Result:
KleidiAI does not include a function specifically named kai_cast_bf16_f32 for BF16 to F32 conversion [1][2][3]. Instead, KleidiAI is primarily a library of performance-critical micro-kernels for operations like matrix multiplication and convolution [1][4]. In the context of KleidiAI, BF16-to-F32 or BF16-based operations are handled through specific operator micro-kernels, particularly those designed for GEMM (General Matrix Multiply) that target BF16 input and F32 output [5][3]. For instance, KleidiAI provides kernels such as matmul_clamp_f32_bf16p_bf16p, which perform matrix multiplication of two BF16 matrices and produce an F32 output [5]. These kernels are integrated into frameworks (like ONNX Runtime or llama.cpp) where the framework handles the necessary data preparation, packing, and invocation of the micro-kernel [6][7][8]. If you are looking to perform raw BF16-to-F32 conversion on Arm CPUs, you would typically use standard Arm instruction set features (such as those available via Neon or SVE/SME) rather than a specific KleidiAI library function, as KleidiAI focuses on optimized computational routines rather than general-purpose data format conversion utilities [1][4]. For example, developers using Arm architectures often utilize VCVT instructions (Neon) or equivalent SVE/SME conversion instructions to handle BF16 and F32 data types [1]. In summary: 1. No direct kai_cast_bf16_f32 API exists in KleidiAI [1][2][3]. 2. BF16/F32 functionality is implemented within dedicated matmul/operator micro-kernels where conversion is implicit in the operation [5][3]. 3. For general-purpose conversion, refer to Arm architecture reference manuals for the relevant conversion instructions (e.g., vcvt for Neon/SVE) [1].
Citations:
- 1: https://github.com/ARM-software/KleidiAI
- 2: https://github.com/ARM-software/kleidiai/blob/main/docs/microkernel_tables.md
- 3: https://github.com/ARM-software/kleidiai/blob/main/CHANGELOG.md
- 4: https://github.com/ARM-software/kleidiai/blob/main/AGENTS.md
- 5: https://github.com/ARM-software/kleidiai/blob/main/docs/README.md
- 6: https://github.com/ggml-org/ggml/blob/c044a8ee/src/ggml-cpu/kleidiai/kernels.cpp
- 7: https://learn.arm.com/tag/kleidiai/
- 8: [MLAS] Integrate KleidiAI BF16 SME2 Kernel Through Mlas SBGEMM Path microsoft/onnxruntime#26773
🌐 Web query:
ARM SVE bf16 f32 to bf16 rounding semantics fcvtns bf16 f32
💡 Result:
In the Arm AArch64 and SVE architectures, there is no instruction named fcvtns for converting floating-point to BFloat16. The instruction commonly used for single-precision (f32) to BFloat16 (bf16) conversion is BFCVT (or BFCVTN) [1][2][3]. The term fcvtns specifically refers to a different class of instructions: Floating-point Convert to Signed integer, with Rounding to Nearest, ties to Even [4][5]. It is not used for floating-point to floating-point conversions like f32 to bf16 [4][6][5]. Regarding BFloat16 conversion semantics: 1. Instruction Selection: For f32 to bf16 conversion, developers use BFCVT (predicated in SVE, scalar in AArch64) [2][3] or BFCVTN (narrowing, vector) [1]. 2. Rounding Semantics: The BFCVT instruction performs the conversion from single-precision to BFloat16 format using the FPConvertBF shared function [1][2][3]. According to official Arm documentation, BFCVT (and the related BFMLAL instructions) honors standard floating-point controls, operating much closer to full IEEE-FP32 accuracy than some other dedicated BF16 matrix-multiply instructions (like BFDOT or BFMMLA, which use a non-IEEE Round to Odd mode) [7]. 3. Contrast with fcvtns: Because fcvtns is an integer conversion instruction, it utilizes the standard FPCR (Floating-point Control Register) rounding modes (including Round to Nearest, ties to even) for converting floating-point values to fixed-point integers [4][6][5]. It is functionally distinct from the BFCVT family of instructions [1][2]. In summary, do not use fcvtns for f32 to bf16 conversion, as it will perform an integer conversion. Use BFCVT or BFCVTN for the intended floating-point narrowing operation [1][2].
Citations:
- 1: https://developer.arm.com/documentation/111108/2026-06/SIMD-FP-Instructions/BFCVTN--BFCVTN2--Single-precision-convert-to-BFloat16--vector--
- 2: https://developer.arm.com/documentation/ddi0602/2026-06/SVE-Instructions/BFCVT--Single-precision-convert-to-BFloat16--predicated--
- 3: https://developer.arm.com/documentation/ddi0602/2026-06/SIMD-FP-Instructions/BFCVT--Single-precision-convert-to-BFloat16--scalar--
- 4: https://developer.arm.com/documentation/ddi0602/2026-06/SIMD-FP-Instructions/FCVTNS--vector---Floating-point-convert-to-signed-integer--rounding-to-nearest-with-ties-to-even--vector--
- 5: https://developer.arm.com/documentation/111108/2026-06/SIMD-FP-Instructions/FCVTNS--scalar---Floating-point-convert-to-signed-integer--rounding-to-nearest-with-ties-to-even--scalar--
- 6: https://developer.arm.com/architectures/instruction-sets/intrinsics/vcvtns_s32_f32
- 7: https://developer.arm.com/community/arm-community-blogs/b/ai-blog/posts/bfloat16-processing-for-neural-networks-on-armv8_2d00_a
Round bf16 scale conversions instead of truncating.
truncateToBFloat16 discards the low 16 bits of the f32 bit pattern, so values like 2.5f dequantify to 4/8 instead of the canonical bf16 encodings 16/32. The scale produced here is written per W4A32 block and passed to KleidiAI with scale_dt = kai_dt_bf16, so use an f32→bf16 conversion with round-to-nearest-even semantics.
🤖 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 `@mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp` around lines 26 - 30,
Update truncateToBFloat16 to convert the float bit pattern to bfloat16 using
round-to-nearest-even rather than simply shifting away the lower 16 bits.
Preserve the returned uint16_t representation and ensure the conversion handles
carry into the upper half correctly before the per-block scale is passed with
kai_dt_bf16.
| KaiW4A32Tile kaiW4A32TileFromName(std::string_view tile_name) { | ||
| if (tile_name == "qai8dxp1x8_qsi4c32p4x8_1x4x32" || tile_name == "qai8dxp4x8_qsi4c32p4x8_8x4x32" | ||
| || tile_name == "qai8dxp4x8_qsi4c32p4x8_16x4x32") { | ||
| return {.nr = 4, .kr = 16, .sr = 2}; | ||
| } | ||
| if (tile_name == "qai8dxp1x8_qsi4c32p8x8_1x8x32" || tile_name == "qai8dxp4x8_qsi4c32p8x8_4x8x32") { | ||
| return {.nr = 8, .kr = 16, .sr = 2}; | ||
| } | ||
| if (tile_name == "qai8dxp1x4_qsi4c32p4x4_1x4") { return {.nr = 4, .kr = 8, .sr = 2}; } | ||
| throw std::invalid_argument("Unsupported KAI W4A32 tile configuration: " + std::string(tile_name)); | ||
| } | ||
|
|
||
| size_t kaiW4A32PackedSize(size_t out_channels, size_t in_channels, std::string_view tile_name) { | ||
| if (out_channels == 0 || in_channels == 0 || (in_channels % kBlockLength) != 0) { | ||
| throw std::invalid_argument("KAI W4A32 requires non-zero dimensions and input channels divisible by 32"); | ||
| } | ||
| const auto tile = kaiW4A32TileFromName(tile_name); | ||
| return kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0(out_channels, in_channels, tile.nr, tile.kr, tile.sr, | ||
| kBlockLength, kai_dt_bf16); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate kai.cpp files =="
fd -a 'kai.*\.cpp|kai.*\.h' . | sed 's#^\./##' | head -100
echo "== locate target file and related includes =="
for f in mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp mllm/backends/cpu/kernels/common/kai.cpp; do
if [ -f "$f" ]; then
echo "--- $f line count ---"
wc -l "$f"
echo "--- includes/outline ---"
ast-grep outline "$f" --view compact || true
echo "--- relevant sections ---"
rg -n "kaiW4a32|W4A32|kai_get_rhs_packed_size_rhs_pack_nxk|get_nr|get_sr|get_cr|tile_name|kai_get_nr_k|kai_get_sr_k|kai_get_cr_k|tile" "$f" -C 3 -A 80 -B 20 | head -360
fi
done
echo "== search for W4A32 tile definitions in repo =="
rg -n "qai8dxp1x8|qai8dxp4x8|qai8dxp1x4|W4A32|kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p|get_rhs_packed_size|nr.*kr.*sr|kai_get_nr|kai_get_sr|kai_get_cr" . -g '*.cpp' -g '*.h' -g '*.hpp' | head -300Repository: UbiquitousLearning/mllm
Length of output: 46720
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate kai.cpp files =="
fd -a 'kai.*\.cpp|kai.*\.h' . | sed 's#^\./##' | head -100
echo "== locate target file and related includes =="
for f in mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp mllm/backends/cpu/kernels/common/kai.cpp; do
if [ -f "$f" ]; then
echo "--- $f line count ---"
wc -l "$f"
echo "--- includes/outline ---"
ast-grep outline "$f" --view compact || true
echo "--- relevant sections ---"
rg -n "kaiW4a32|W4A32|kai_get_rhs_packed_size_rhs_pack_nxk|get_nr|get_sr|get_cr|tile_name|kai_get_nr_k|kai_get_sr_k|kai_get_cr_k|tile" "$f" -C 3 -A 120 -B 40 | head -500
fi
done
echo "== search for W4A32 tile definitions in repo =="
rg -n "qai8dxp1x8|qai8dxp4x8|qai8dxp1x4|W4A32|kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p|get_rhs_packed_size|nr.*kr.*sr|kai_get_nr|kai_get_sr|kai_get_cr" . -g '*.cpp' -g '*.h' -g '*.hpp' | head -300Repository: UbiquitousLearning/mllm
Length of output: 47526
Add an ARM cross-check for W4A32 tile parameters.
kaiW4A32TileFromName hardcodes nr/kr/sr for all supported W4A32 tiles, while the ARM ukernels expose get_nr/get_kr/get_sr for the same tile names. The existing non-ARM determinism test only exercises one tile; add an ARM-only check over every tile so the packed-size/packing parameters cannot drift from the real KAI microkernel values after a KleidiAI update.
🤖 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 `@mllm/backends/cpu/kernels/common/kai_w4a32_pack.cpp` around lines 78 - 97,
Add an ARM-only test that iterates over every tile name supported by
kaiW4A32TileFromName and compares its returned nr, kr, and sr values with the
corresponding ARM KAI ukernel get_nr, get_kr, and get_sr values. Reuse the
existing tile-to-ukernel mapping and test infrastructure, and retain the current
non-ARM determinism test unchanged.
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with 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.
Inline comments:
In `@mllm/models/qwen3_5/tokenization_qwen3_5.hpp`:
- Around line 19-28: Document the public API for Qwen3_5StreamingUtf8Decoder and
Qwen3_5Tokenizer, including each public constructor and method in the referenced
ranges. Describe class purpose, parameter and return semantics,
streaming/replacement behavior, and every runtime_error condition; ensure
comments accurately distinguish append, finish, reset, and tokenizer operations.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 37c0d66e-8397-4489-b89f-51a29f28906f
📒 Files selected for processing (7)
examples/qwen3_5/main.cppmllm/models/qwen3_5/configuration_qwen3_5.hppmllm/models/qwen3_5/tokenization_qwen3_5.hpppymllm/mobile/tests/test_convertor.pytests/cpu/CMakeLists.txttests/cpu/Qwen35ConfigTest.cpptests/cpu/Qwen35TokenizerTest.cpp
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/qwen3_5/main.cpp
- mllm/models/qwen3_5/configuration_qwen3_5.hpp
| class Qwen3_5StreamingUtf8Decoder { | ||
| public: | ||
| std::string append(std::string_view bytes) { | ||
| if (!bytes.empty()) { pending_.append(bytes.data(), bytes.size()); } | ||
| return drain(false); | ||
| } | ||
|
|
||
| std::string finish() { return drain(true); } | ||
|
|
||
| void reset() { pending_.clear(); } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the public tokenizer APIs.
Qwen3_5StreamingUtf8Decoder and Qwen3_5Tokenizer expose undocumented classes and methods, including replacement behavior and runtime_error paths. Add API comments covering purpose, parameters, returns, and errors. As per coding guidelines, “Ensure public APIs, classes, and functions have clear docstrings or comments explaining purpose, parameters, returns, and errors.”
Also applies to: 242-250, 284-367
🤖 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 `@mllm/models/qwen3_5/tokenization_qwen3_5.hpp` around lines 19 - 28, Document
the public API for Qwen3_5StreamingUtf8Decoder and Qwen3_5Tokenizer, including
each public constructor and method in the referenced ranges. Describe class
purpose, parameter and return semantics, streaming/replacement behavior, and
every runtime_error condition; ensure comments accurately distinguish append,
finish, reset, and tokenizer operations.
Source: Coding guidelines
2 participants
What this PR does
Adds an end-to-end, text-only Qwen3.5-0.8B path for the CPU backend:
(batch, value_head)parallelismThe vision tower and MTP layer are deliberately excluded.
Review guide
mllm/models/qwen3_5/gated_delta_net.*,kai_w4a32_pack.*,ContiguousOp.cpppymllm/mobile/convertor/,pymllm/mobile/quantize/examples/qwen3_5/,tests/cpu/,tests/core/Suggested review order: model contract → GDN kernel → W4A32 conversion → runner/tests.
Current-head validation
Exact PR HEAD:
6cf4ea1eValidated implementation HEAD:
0576fbde(the later commits only update the root README model table).git diff --checkpassedmllm-qwen3-5-runner --helppassedEarlier H20, Android, and device runs predate the validated implementation HEAD and are therefore retained below as historical evidence, not presented as current-head validation.
Earlier cross-platform, correctness, and performance evidence
896bfd45e42ce170Hk=16,Hv=32)e2492f221.2684x, 21.16% lower)e2492f22through current HEAD6cf4ea1e.How to use it
examples/qwen3_5/README.mddocuments checkpoint validation, text-tower conversion, desktop/Android builds, and runner invocation.Known limits
Hk < Hvis covered by a focused test with official 4B/9B geometry, not by an end-to-end performance claimTracks #651.
Related context: #657. This implementation and its validation were developed independently; #657 was not used as a code baseline.
Summary by CodeRabbit